""" 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()