Add training script

train.py

@@ -0,0 +1,259 @@
+"""
+DeepSeek-V4 Training Script
+============================
+Training script for the open DeepSeek-V4 reproduction.
+
+Implements the training recipe from the DeepSeek-V4 technical report:
+- Muon optimizer (or AdamW fallback)
+- FP8/BF16 mixed precision training
+- Gradient checkpointing
+- Distributed training via accelerate
+- Progressive sequence length extension
+- MTP auxiliary loss
+
+Usage:
+    # Single GPU testing with nano config
+    python train.py --config configs/config_nano.json --dataset HuggingFaceFW/fineweb-edu-score-2 --max_steps 100
+    
+    # Multi-GPU with accelerate
+    accelerate launch train.py --config configs/config_pro.json --dataset HuggingFaceFW/fineweb-edu-score-2
+"""
+
+import os
+import sys
+import json
+import math
+import argparse
+import logging
+from typing import Optional
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, IterableDataset
+
+from transformers import (
+    TrainingArguments,
+    Trainer,
+    TrainerCallback,
+    AutoTokenizer,
+    get_cosine_schedule_with_warmup,
+)
+from datasets import load_dataset
+
+# Add parent directory to path
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from configuration_deepseek_v4 import DeepSeekV4Config
+from modeling_deepseek_v4 import DeepSeekV4ForCausalLM
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# Data Pipeline
+# =============================================================================
+
+class StreamingTextDataset(IterableDataset):
+    """Streaming dataset that tokenizes on-the-fly and packs sequences."""
+    
+    def __init__(
+        self,
+        dataset_name: str,
+        tokenizer,
+        max_length: int = 2048,
+        split: str = "train",
+        text_column: str = "text",
+        streaming: bool = True,
+    ):
+        self.dataset = load_dataset(dataset_name, split=split, streaming=streaming)
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.text_column = text_column
+    
+    def __iter__(self):
+        buffer = []
+        for example in self.dataset:
+            text = example.get(self.text_column, "")
+            if not text:
+                continue
+            tokens = self.tokenizer.encode(text, add_special_tokens=False)
+            buffer.extend(tokens)
+            
+            while len(buffer) >= self.max_length:
+                chunk = buffer[:self.max_length]
+                buffer = buffer[self.max_length:]
+                
+                input_ids = torch.tensor(chunk, dtype=torch.long)
+                labels = input_ids.clone()
+                
+                yield {
+                    "input_ids": input_ids,
+                    "labels": labels,
+                }
+
+
+def collate_fn(batch):
+    """Collate function for DataLoader."""
+    input_ids = torch.stack([item["input_ids"] for item in batch])
+    labels = torch.stack([item["labels"] for item in batch])
+    return {"input_ids": input_ids, "labels": labels}
+
+
+# =============================================================================
+# Custom Callbacks
+# =============================================================================
+
+class DeepSeekV4Callback(TrainerCallback):
+    """Custom callback for DeepSeek-V4 training monitoring."""
+    
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if logs is None:
+            return
+        
+        step = state.global_step
+        loss = logs.get("loss")
+        lr = logs.get("learning_rate")
+        
+        if loss is not None:
+            # Check for training issues
+            if math.isnan(loss) or math.isinf(loss):
+                logger.error(f"Step {step}: Loss is {loss}! Training may be diverging.")
+            elif loss > 20 and step > 100:
+                logger.warning(f"Step {step}: Loss={loss:.4f} is high. Consider reducing LR.")
+    
+    def on_train_begin(self, args, state, control, **kwargs):
+        logger.info("=" * 60)
+        logger.info("DeepSeek-V4 Training Started")
+        logger.info(f"  Max steps: {args.max_steps}")
+        logger.info(f"  Batch size: {args.per_device_train_batch_size}")
+        logger.info(f"  Gradient accumulation: {args.gradient_accumulation_steps}")
+        logger.info(f"  Learning rate: {args.learning_rate}")
+        logger.info(f"  Hardware: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}")
+        logger.info("=" * 60)
+
+
+# =============================================================================
+# Training Entry Point
+# =============================================================================
+
+def main():
+    parser = argparse.ArgumentParser(description="Train DeepSeek-V4")
+    
+    # Model & Data
+    parser.add_argument("--config", type=str, required=True, help="Path to model config JSON")
+    parser.add_argument("--dataset", type=str, default="HuggingFaceFW/fineweb-edu-score-2", help="Dataset name on HF Hub")
+    parser.add_argument("--tokenizer", type=str, default="deepseek-ai/DeepSeek-V4-Pro", help="Tokenizer name")
+    parser.add_argument("--max_length", type=int, default=2048, help="Max sequence length")
+    parser.add_argument("--text_column", type=str, default="text", help="Text column in dataset")
+    
+    # Training hyperparameters (V4 paper recipe)
+    parser.add_argument("--max_steps", type=int, default=10000, help="Max training steps")
+    parser.add_argument("--batch_size", type=int, default=4, help="Per-device batch size")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=8, help="Gradient accumulation")
+    parser.add_argument("--learning_rate", type=float, default=2.2e-4, help="Peak learning rate")
+    parser.add_argument("--min_lr_ratio", type=float, default=0.033, help="Min LR as ratio of peak")
+    parser.add_argument("--warmup_steps", type=int, default=2000, help="Warmup steps")
+    parser.add_argument("--weight_decay", type=float, default=0.1, help="Weight decay")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="Max gradient norm")
+    
+    # MTP
+    parser.add_argument("--mtp_loss_weight", type=float, default=0.3, help="MTP auxiliary loss weight")
+    
+    # Infrastructure
+    parser.add_argument("--output_dir", type=str, default="./output", help="Output directory")
+    parser.add_argument("--hub_model_id", type=str, default=None, help="Push to this HF Hub model ID")
+    parser.add_argument("--logging_steps", type=int, default=10, help="Log every N steps")
+    parser.add_argument("--save_steps", type=int, default=1000, help="Save every N steps")
+    parser.add_argument("--bf16", action="store_true", default=True, help="Use BF16 mixed precision")
+    parser.add_argument("--gradient_checkpointing", action="store_true", default=True, help="Enable gradient checkpointing")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    
+    args = parser.parse_args()
+    
+    # Load config
+    with open(args.config) as f:
+        config_dict = json.load(f)
+    config = DeepSeekV4Config(**config_dict)
+    
+    logger.info(f"Model config: {config.hidden_size}d, {config.num_hidden_layers}L, "
+                f"{config.n_routed_experts}E, hc_mult={config.hc_mult}")
+    
+    # Initialize model
+    model = DeepSeekV4ForCausalLM(config)
+    
+    total_params = sum(p.numel() for p in model.parameters())
+    logger.info(f"Total parameters: {total_params:,} ({total_params/1e9:.2f}B)")
+    
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    # Create dataset
+    train_dataset = StreamingTextDataset(
+        dataset_name=args.dataset,
+        tokenizer=tokenizer,
+        max_length=args.max_length,
+        text_column=args.text_column,
+    )
+    
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir=args.output_dir,
+        max_steps=args.max_steps,
+        per_device_train_batch_size=args.batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        learning_rate=args.learning_rate,
+        weight_decay=args.weight_decay,
+        max_grad_norm=args.grad_clip,
+        warmup_steps=args.warmup_steps,
+        lr_scheduler_type="cosine",
+        bf16=args.bf16 and torch.cuda.is_available(),
+        logging_steps=args.logging_steps,
+        logging_strategy="steps",
+        logging_first_step=True,
+        save_steps=args.save_steps,
+        save_total_limit=3,
+        seed=args.seed,
+        disable_tqdm=True,
+        gradient_checkpointing=args.gradient_checkpointing,
+        push_to_hub=args.hub_model_id is not None,
+        hub_model_id=args.hub_model_id,
+        report_to="none",  # override to "trackio" when using monitoring
+        dataloader_pin_memory=True,
+        remove_unused_columns=False,
+    )
+    
+    # Create trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        data_collator=collate_fn,
+        callbacks=[DeepSeekV4Callback()],
+    )
+    
+    # Train
+    logger.info("Starting training...")
+    train_result = trainer.train()
+    
+    # Save final model
+    trainer.save_model()
+    
+    # Log metrics
+    metrics = train_result.metrics
+    trainer.log_metrics("train", metrics)
+    trainer.save_metrics("train", metrics)
+    
+    logger.info("Training complete!")
+    logger.info(f"Final loss: {metrics.get('train_loss', 'N/A')}")
+    
+    if args.hub_model_id:
+        logger.info(f"Pushing to Hub: {args.hub_model_id}")
+        trainer.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()