Add pretrain.py — pretrain_domain_model with HF Trainer, cosine schedule, DataCollatorForLanguageModeling

src/domain_tokenizer/training/pretrain.py

@@ -0,0 +1,132 @@
+"""
+Pre-training function for DomainTransformer.
+
+Uses HuggingFace Trainer with DataCollatorForLanguageModeling(mlm=False)
+which automatically sets labels = input_ids and masks padding with -100.
+
+Usage:
+    from domain_tokenizer.training import pretrain_domain_model, prepare_clm_dataset
+    dataset = prepare_clm_dataset(user_sequences, builder, hf_tokenizer, block_size=512)
+    config = DomainTransformerConfig.from_preset("24m", vocab_size=hf_tokenizer.vocab_size)
+    model = DomainTransformerForCausalLM(config)
+    pretrain_domain_model(model, hf_tokenizer, dataset)
+"""
+
+import logging
+from typing import Optional
+
+from datasets import Dataset as HFDataset
+from transformers import (
+    DataCollatorForLanguageModeling,
+    PreTrainedTokenizerFast,
+    Trainer,
+    TrainingArguments,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def pretrain_domain_model(
+    model,
+    tokenizer: PreTrainedTokenizerFast,
+    train_dataset: HFDataset,
+    eval_dataset: Optional[HFDataset] = None,
+    output_dir: str = "./domain_pretrain_checkpoints",
+    hub_model_id: Optional[str] = None,
+    num_epochs: int = 10,
+    per_device_batch_size: int = 32,
+    gradient_accumulation_steps: int = 4,
+    learning_rate: float = 3e-4,
+    lr_scheduler_type: str = "cosine",
+    warmup_steps: int = 500,
+    weight_decay: float = 0.01,
+    max_grad_norm: float = 1.0,
+    bf16: bool = False,
+    fp16: bool = False,
+    logging_steps: int = 50,
+    save_steps: int = 500,
+    eval_steps: int = 500,
+    save_total_limit: int = 3,
+    dataloader_num_workers: int = 4,
+    report_to: str = "none",
+    run_name: Optional[str] = None,
+    seed: int = 42,
+    gradient_checkpointing: bool = False,
+    resume_from_checkpoint: Optional[str] = None,
+    **extra_training_args,
+) -> Trainer:
+    """Pre-train a DomainTransformerForCausalLM with HF Trainer.
+
+    The dataset should be packed via prepare_clm_dataset() for 100% token utilization.
+
+    Returns:
+        The Trainer instance (for inspection, continued training, etc.).
+    """
+    if tokenizer.pad_token_id is None:
+        raise ValueError(
+            "Tokenizer must have pad_token set. "
+            "DomainTokenizerBuilder.build() should set this automatically."
+        )
+
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+    push_to_hub = hub_model_id is not None
+
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        num_train_epochs=num_epochs,
+        per_device_train_batch_size=per_device_batch_size,
+        per_device_eval_batch_size=per_device_batch_size,
+        gradient_accumulation_steps=gradient_accumulation_steps,
+        learning_rate=learning_rate,
+        lr_scheduler_type=lr_scheduler_type,
+        warmup_steps=warmup_steps,
+        weight_decay=weight_decay,
+        max_grad_norm=max_grad_norm,
+        bf16=bf16, fp16=fp16,
+        logging_strategy="steps",
+        logging_steps=logging_steps,
+        logging_first_step=True,
+        disable_tqdm=True,
+        eval_strategy="steps" if eval_dataset else "no",
+        eval_steps=eval_steps if eval_dataset else None,
+        save_strategy="steps",
+        save_steps=save_steps,
+        save_total_limit=save_total_limit,
+        push_to_hub=push_to_hub,
+        hub_model_id=hub_model_id if push_to_hub else None,
+        dataloader_num_workers=dataloader_num_workers,
+        report_to=report_to,
+        run_name=run_name,
+        seed=seed,
+        gradient_checkpointing=gradient_checkpointing,
+        remove_unused_columns=True,
+        **extra_training_args,
+    )
+
+    effective_batch = per_device_batch_size * gradient_accumulation_steps
+    n_params = sum(p.numel() for p in model.parameters())
+
+    logger.info(f"=== Domain Pre-Training ===")
+    logger.info(f"  Model params: {n_params:,}")
+    logger.info(f"  Train samples: {len(train_dataset):,}")
+    logger.info(f"  Block size: {len(train_dataset[0]['input_ids'])}")
+    logger.info(f"  Batch size: {per_device_batch_size} x {gradient_accumulation_steps} = {effective_batch}")
+    logger.info(f"  Epochs: {num_epochs}, LR: {learning_rate} ({lr_scheduler_type})")
+    logger.info(f"  Push to hub: {hub_model_id if push_to_hub else 'disabled'}")
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        data_collator=data_collator,
+        processing_class=tokenizer,
+    )
+
+    trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+
+    if push_to_hub:
+        logger.info(f"Pushing model to hub: {hub_model_id}")
+        trainer.push_to_hub()
+
+    return trainer