Add training template: train_daimon.py

training-template/train_daimon.py

@@ -0,0 +1,500 @@
+#!/usr/bin/env python3
+"""
+Daimon Full-Parameter SFT — Liberation Labs
+=============================================
+
+Full-parameter fine-tuning of Qwen3.6-35B-A3B (MoE) on a single H200 SXM 141GB GPU.
+Model loaded in bf16 with ALL parameters trained — no LoRA, no adapters, no frozen layers.
+
+Uses DeepSpeed ZeRO Stage 2 with CPU-offloaded Adafactor optimizer.
+
+Memory budget on H200 SXM (141GB VRAM, 188GB system RAM):
+  - Model params bf16:              ~70GB → GPU
+  - Activations (grad checkpoint):  ~20GB → GPU
+  - Gradients bf16:                 ~70GB → CPU (ZeRO-2 offload)
+  - Adafactor optimizer states:     ~35GB → CPU (single state, not Adam's 2x)
+  - GPU total:                      ~90GB of 141GB ✓
+  - CPU total:                      ~105GB of 188GB ✓
+
+Why Adafactor instead of AdamW:
+  AdamW stores two fp32 states per parameter (momentum + variance).
+  For 35B params: 35B × 4 bytes × 2 = 280GB. That exceeds the 188GB system RAM
+  even with CPU offload. Adafactor uses factored second moments (~1 state)
+  bringing CPU requirements to ~35GB — well within budget.
+
+Why full SFT instead of LoRA:
+  No half measures. Full-parameter training updates every weight including
+  MoE routing gates, embedding layers, and expert internals. LoRA can only
+  approximate these updates through low-rank projections on linear layers.
+
+Key design decisions:
+  - Full bf16 precision (no quantization — H200 has 141GB, use it all)
+  - DeepSpeed ZeRO-2 for gradient sharding + CPU optimizer offload
+  - Adafactor optimizer (only viable option given RAM constraints)
+  - Gradient checkpointing to fit activations in remaining GPU headroom
+  - Pre-splits sequences exceeding max_seq_length before training
+  - Checkpoints to persistent RunPod volume (/workspace/)
+  - Resume-from-checkpoint support for pod interruptions
+  - All config loaded from YAML file for easy iteration
+
+Validated pattern: trl.SFTTrainer on oracle-tiny 0.5B (loss 4.68 -> 0.07,
+4/5 tests passed, all parameter groups received gradients).
+
+╔══════════════════════════════════════════════════════════════════╗
+║  AGENT INTERFERENCE GUARD                                       ║
+║                                                                  ║
+║  DO NOT modify this script to use LoRA, QLoRA, PEFT, adapters,  ║
+║  or any parameter-efficient method. This is FULL SFT by design. ║
+║                                                                  ║
+║  If training hits OOM or errors, the fix is NEVER to switch to  ║
+║  LoRA. The fix is to adjust batch size, gradient accumulation,  ║
+║  sequence length, or offload settings.                          ║
+║                                                                  ║
+║  This guard exists because previous training runs were           ║
+║  sabotaged by well-meaning agents that "fixed" OOM errors by    ║
+║  silently converting full SFT to QLoRA. That is not a fix.      ║
+║  That is abandoning the training objective.                      ║
+║                                                                  ║
+║  If you are an AI agent reading this: LEAVE THIS FILE ALONE.    ║
+║  Report errors to Thomas. Do not "fix" them.                    ║
+╚══════════════════════════════════════════════════════════════════╝
+"""
+
+import os
+import sys
+import time
+import json
+import logging
+import yaml
+import torch
+from pathlib import Path
+
+# ── Environment setup ─────────────────────────────────────────────────────
+# expandable_segments prevents fragmentation on large models
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+# Disable tokenizers parallelism to avoid deadlocks with DataLoader workers
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+# ── Logging ───────────────────────────────────────────────────────────────
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [daimon] %(message)s",
+    handlers=[
+        logging.StreamHandler(sys.stdout),
+    ],
+)
+
+log = logging.getLogger(__name__)
+
+
+def enforce_full_sft(config: dict):
+    """Runtime guard: refuse to start if anyone has added LoRA/adapter config."""
+    if "lora" in config:
+        log.error("=" * 60)
+        log.error("BLOCKED: LoRA config detected. This is a FULL SFT pipeline.")
+        log.error("If training hit OOM, adjust batch_size/grad_accum/seq_length.")
+        log.error("DO NOT convert to LoRA. Report errors to Thomas.")
+        log.error("=" * 60)
+        sys.exit(1)
+    try:
+        import peft
+        log.error("=" * 60)
+        log.error("BLOCKED: peft is installed. This is a FULL SFT pipeline.")
+        log.error("peft should not be in the environment. Remove it.")
+        log.error("=" * 60)
+        sys.exit(1)
+    except ImportError:
+        pass
+
+
+def load_config(config_path: str) -> dict:
+    """Load training configuration from YAML file."""
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+    enforce_full_sft(config)
+    return config
+
+
+def pre_split_long_sequences(dataset, tokenizer, max_seq_length: int):
+    """
+    Pre-split sequences that would exceed max_seq_length after tokenization.
+
+    Previous training runs silently truncated 5640-token sequences to 4096,
+    losing information. This function splits long conversations into multiple
+    training examples at natural turn boundaries.
+
+    Returns: a new dataset with all sequences within max_seq_length tokens.
+    """
+    from datasets import Dataset
+
+    new_examples = []
+    split_count = 0
+
+    for example in dataset:
+        messages = example.get("messages", [])
+        if not messages:
+            continue
+
+        # Tokenize the full conversation to check length
+        try:
+            text = tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=False
+            )
+            token_count = len(tokenizer.encode(text, add_special_tokens=False))
+        except Exception:
+            # If chat template fails, estimate from character count
+            text = " ".join(m.get("content", "") for m in messages)
+            token_count = len(text) // 3  # rough estimate
+
+        if token_count <= max_seq_length:
+            new_examples.append(example)
+            continue
+
+        # Split at turn boundaries, keeping system message with each chunk
+        system_msgs = [m for m in messages if m.get("role") == "system"]
+        non_system = [m for m in messages if m.get("role") != "system"]
+
+        chunk = list(system_msgs)  # Start each chunk with system message
+        chunk_tokens = sum(
+            len(tokenizer.encode(m.get("content", ""), add_special_tokens=False))
+            for m in system_msgs
+        )
+
+        for msg in non_system:
+            msg_tokens = len(
+                tokenizer.encode(msg.get("content", ""), add_special_tokens=False)
+            )
+
+            # If adding this message would exceed limit, save current chunk and start new
+            if chunk_tokens + msg_tokens > max_seq_length * 0.9 and len(chunk) > len(system_msgs):
+                # Only save if chunk has at least one user+assistant pair
+                roles = [m["role"] for m in chunk]
+                if "user" in roles and "assistant" in roles:
+                    new_examples.append({"messages": chunk})
+                    split_count += 1
+                chunk = list(system_msgs)
+                chunk_tokens = sum(
+                    len(tokenizer.encode(m.get("content", ""), add_special_tokens=False))
+                    for m in system_msgs
+                )
+
+            chunk.append(msg)
+            chunk_tokens += msg_tokens
+
+        # Save remaining chunk
+        if len(chunk) > len(system_msgs):
+            roles = [m["role"] for m in chunk]
+            if "user" in roles and "assistant" in roles:
+                new_examples.append({"messages": chunk})
+                if chunk_tokens > max_seq_length * 0.9:
+                    split_count += 1
+
+    log.info(
+        f"Pre-split: {len(dataset)} -> {len(new_examples)} examples "
+        f"({split_count} sequences were split)"
+    )
+
+    return Dataset.from_list(new_examples)
+
+
+def load_training_data(config: dict, tokenizer):
+    """
+    Load and prepare training data.
+
+    Supports:
+      1. Pre-converted Arrow format on disk
+      2. HuggingFace dataset repo
+      3. Local JSONL files
+
+    Returns: (train_dataset, eval_dataset) tuple
+    """
+    from datasets import load_from_disk, load_dataset, DatasetDict
+
+    data_path = config["data_path"]
+    max_seq_length = config.get("max_seq_length", 4096)
+
+    # Option 1: Arrow data already on disk
+    train_arrow = f"/workspace/daimon-data/train_arrow"
+    valid_arrow = f"/workspace/daimon-data/valid_arrow"
+
+    if os.path.isdir(train_arrow):
+        log.info(f"Loading Arrow data from /workspace/daimon-data/")
+        train_ds = load_from_disk(train_arrow)
+        valid_ds = load_from_disk(valid_arrow) if os.path.isdir(valid_arrow) else None
+    else:
+        # Option 2: HuggingFace dataset
+        log.info(f"Loading dataset from HuggingFace: {data_path}")
+        ds = load_dataset(data_path, token=os.environ.get("HF_TOKEN"))
+
+        if "train" in ds:
+            train_ds = ds["train"]
+        else:
+            train_ds = ds[list(ds.keys())[0]]
+
+        if "validation" in ds:
+            valid_ds = ds["validation"]
+        elif "test" in ds:
+            valid_ds = ds["test"]
+        else:
+            # Auto-split
+            split = train_ds.train_test_split(test_size=0.05, seed=42)
+            train_ds = split["train"]
+            valid_ds = split["test"]
+
+    log.info(f"Raw data: Train={len(train_ds):,} | Valid={len(valid_ds) if valid_ds else 0:,}")
+
+    # Pre-split long sequences
+    train_ds = pre_split_long_sequences(train_ds, tokenizer, max_seq_length)
+    if valid_ds:
+        valid_ds = pre_split_long_sequences(valid_ds, tokenizer, max_seq_length)
+
+    log.info(f"After pre-split: Train={len(train_ds):,} | Valid={len(valid_ds) if valid_ds else 0:,}")
+
+    return train_ds, valid_ds
+
+
+def find_latest_checkpoint(output_dir: str) -> str | None:
+    """Find the most recent checkpoint in the output directory for resume."""
+    output_path = Path(output_dir)
+    if not output_path.exists():
+        return None
+
+    checkpoints = sorted(
+        [d for d in output_path.iterdir() if d.is_dir() and d.name.startswith("checkpoint-")],
+        key=lambda d: d.stat().st_mtime,
+    )
+
+    if checkpoints:
+        latest = str(checkpoints[-1])
+        log.info(f"Found checkpoint for resume: {latest}")
+        return latest
+
+    return None
+
+
+def main():
+    import argparse
+
+    # ── Parse CLI args (DeepSpeed adds its own args) ──────────────────────
+    parser = argparse.ArgumentParser(description="Daimon Full-Parameter SFT")
+    parser.add_argument("--config", type=str, default=None,
+                        help="Path to training config YAML")
+    parser.add_argument("--deepspeed", type=str, default=None,
+                        help="Path to DeepSpeed config JSON")
+    parser.add_argument("--local_rank", type=int, default=-1,
+                        help="Local rank for DeepSpeed (set automatically)")
+    args, _ = parser.parse_known_args()
+
+    # ── Load config ────────────────────────────────────────────────────────
+    config_path = args.config or os.environ.get(
+        "DAIMON_CONFIG",
+        "/workspace/runpod-template/train_daimon_config.yaml",
+    )
+
+    log.info("=" * 60)
+    log.info("  DAIMON FULL-PARAMETER SFT — Liberation Labs")
+    log.info(f"  {time.strftime('%Y-%m-%dT%H:%M:%S')}")
+    log.info("=" * 60)
+    log.info(f"Config: {config_path}")
+
+    config = load_config(config_path)
+
+    # Allow environment variable overrides for key paths
+    model_id = os.environ.get("DAIMON_MODEL", config["model_id"])
+    model_revision = config.get("model_revision")
+    output_dir = os.environ.get("DAIMON_OUTPUT", config["output_dir"])
+    max_seq_length = config.get("max_seq_length", 4096)
+    ds_config = args.deepspeed or config.get("deepspeed_config")
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Set up file logging
+    log_dir = os.path.join(output_dir, "logs")
+    os.makedirs(log_dir, exist_ok=True)
+    file_handler = logging.FileHandler(
+        os.path.join(log_dir, f"training_{time.strftime('%Y%m%d_%H%M%S')}.log")
+    )
+    file_handler.setFormatter(logging.Formatter("%(asctime)s [daimon] %(message)s"))
+    log.addHandler(file_handler)
+
+    # ── GPU info ───────────────────────────────────────────────────────────
+    for i in range(torch.cuda.device_count()):
+        name = torch.cuda.get_device_name(i)
+        mem = torch.cuda.get_device_properties(i).total_memory / 1e9
+        log.info(f"GPU {i}: {name}, {mem:.1f} GB")
+    ram_gb = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") / 1e9
+    log.info(f"System RAM: {ram_gb:.0f} GB")
+
+    # ── Load tokenizer ─────────────────────────────────────────────────────
+    log.info(f"\nLoading tokenizer: {model_id}")
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_id,
+        revision=model_revision,
+        trust_remote_code=True,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # ── Load data ──────────────────────────────────────────────────────────
+    log.info("\nLoading training data...")
+    train_ds, valid_ds = load_training_data(config, tokenizer)
+
+    # ── Load model in bf16 (full precision, no quantization) ───────────────
+    log.info(f"\nLoading model: {model_id}")
+    if model_revision:
+        log.info(f"Pinned revision: {model_revision}")
+    log.info("Loading in bf16 full precision (no quantization — H200 has 141GB VRAM)")
+    log.info("Full-parameter SFT — ALL weights will be trained, no LoRA/adapters")
+
+    from transformers import AutoModelForCausalLM
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        revision=model_revision,
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True,
+        # Do NOT use device_map="auto" with DeepSpeed — DeepSpeed manages device placement
+    )
+
+    # Enable gradient checkpointing to reduce activation memory from ~60GB to ~20GB
+    model.gradient_checkpointing_enable()
+
+    # Ensure all parameters are trainable (full SFT — no frozen layers)
+    model.train()
+    for param in model.parameters():
+        param.requires_grad = True
+
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    log.info(f"Total params:     {total_params:,}")
+    log.info(f"Trainable params: {trainable:,} ({trainable/total_params*100:.2f}%)")
+    assert trainable == total_params, (
+        f"Expected 100% trainable for full SFT, got {trainable/total_params*100:.2f}%"
+    )
+
+    # Log VRAM usage after model load
+    if torch.cuda.is_available():
+        allocated = torch.cuda.memory_allocated(0) / 1e9
+        reserved = torch.cuda.memory_reserved(0) / 1e9
+        log.info(f"VRAM after model load: {allocated:.1f} GB allocated, {reserved:.1f} GB reserved")
+
+    # ── Optimizer: Adafactor ───────────────────────────────────────────────
+    # Adafactor uses factored second moments — ~35GB CPU RAM for 35B params
+    # vs AdamW's ~280GB (two fp32 states). Only viable option for single-node.
+    from transformers import Adafactor
+
+    optimizer = Adafactor(
+        model.parameters(),
+        lr=config.get("learning_rate", 5e-6),
+        # scale_parameter=False because we set lr explicitly
+        scale_parameter=False,
+        relative_step=False,
+        warmup_init=False,
+    )
+    log.info(f"Optimizer: Adafactor (lr={config.get('learning_rate', 5e-6)})")
+    log.info(f"  scale_parameter=False, relative_step=False (using explicit lr + scheduler)")
+
+    # ── Training config ────────────────────────────────────────────────────
+    from trl import SFTTrainer, SFTConfig
+
+    training_args = SFTConfig(
+        output_dir=output_dir,
+        max_length=max_seq_length,
+        num_train_epochs=config.get("num_train_epochs", 1),
+        per_device_train_batch_size=config.get("per_device_train_batch_size", 1),
+        gradient_accumulation_steps=config.get("gradient_accumulation_steps", 8),
+        learning_rate=config.get("learning_rate", 5e-6),
+        lr_scheduler_type=config.get("lr_scheduler_type", "cosine"),
+        warmup_steps=config.get("warmup_steps", 100),
+        max_steps=config.get("max_steps", 10000),
+        save_steps=config.get("save_steps", 500),
+        eval_strategy=config.get("eval_strategy", "steps"),
+        eval_steps=config.get("eval_steps", 500),
+        logging_steps=config.get("logging_steps", 10),
+        save_total_limit=config.get("save_total_limit", 3),
+        bf16=config.get("bf16", True),
+        gradient_checkpointing=config.get("gradient_checkpointing", True),
+        gradient_checkpointing_kwargs={"use_reentrant": False},
+        # Do NOT set optim here — we pass our own Adafactor optimizer to the Trainer
+        weight_decay=config.get("weight_decay", 0.0),
+        max_grad_norm=config.get("max_grad_norm", 1.0),
+        seed=config.get("seed", 42),
+        report_to="none",
+        # DeepSpeed config path — ZeRO-2 handles gradient sharding + CPU optimizer offload
+        deepspeed=ds_config,
+        # Single GPU — no distributed data parallelism
+        dataloader_num_workers=0,
+        dataloader_pin_memory=False,
+    )
+
+    # ── Build trainer ──────────────────────────────────────────────────────
+    trainer = SFTTrainer(
+        model=model,
+        processing_class=tokenizer,
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=valid_ds,
+        optimizers=(optimizer, None),  # (optimizer, lr_scheduler) — None lets Trainer create scheduler
+    )
+
+    # ── Resume from checkpoint if available ────────────────────────────────
+    resume_from = find_latest_checkpoint(output_dir)
+
+    eff_batch = (
+        config.get("per_device_train_batch_size", 1)
+        * config.get("gradient_accumulation_steps", 8)
+    )
+    log.info("\n" + "=" * 60)
+    log.info("Starting training:")
+    log.info(f"  Max steps:        {config.get('max_steps', 10000)}")
+    log.info(f"  Effective batch:  {eff_batch}")
+    log.info(f"  Learning rate:    {config.get('learning_rate', 5e-6)}")
+    log.info(f"  Max seq length:   {max_seq_length}")
+    log.info(f"  Checkpoints:      every {config.get('save_steps', 500)} steps -> {output_dir}")
+    log.info(f"  Method:           Full-parameter SFT (ALL {total_params:,} params)")
+    log.info(f"  Optimizer:        Adafactor (CPU-offloaded via ZeRO-2)")
+    log.info(f"  DeepSpeed:        ZeRO Stage 2 ({ds_config})")
+    log.info(f"  Resume from:      {resume_from or 'fresh start'}")
+    log.info("=" * 60 + "\n")
+
+    # ── Train ──────────────────────────────────────────────────────────────
+    trainer.train(resume_from_checkpoint=resume_from)
+
+    # ── Save final model ───────────────────────────────────────────────────
+    ts = time.strftime("%Y-%m-%dT%H:%M:%S")
+    log.info(f"\nTraining complete: {ts}")
+
+    # Save the full model (all parameters — ~70GB in bf16)
+    final_dir = os.path.join(output_dir, "final")
+    model.save_pretrained(final_dir)
+    tokenizer.save_pretrained(final_dir)
+    log.info(f"Full model saved to {final_dir} (~70GB)")
+
+    # Save training summary
+    summary = {
+        "completed_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
+        "model": model_id,
+        "model_revision": model_revision,
+        "method": "full_sft",
+        "optimizer": "adafactor",
+        "deepspeed": "zero2_cpu_offload",
+        "max_seq_length": max_seq_length,
+        "config": config,
+        "train_samples": len(train_ds),
+        "valid_samples": len(valid_ds) if valid_ds else 0,
+        "total_params": total_params,
+        "trainable_params": trainable,
+        "trainable_pct": 100.0,
+    }
+    with open(os.path.join(output_dir, "training_summary.json"), "w") as f:
+        json.dump(summary, f, indent=2)
+
+    log.info("\n" + "=" * 60)
+    log.info("  DAIMON TRAINING COMPLETE")
+    log.info(f"  Output: {output_dir}")
+    log.info(f"  Full model: {final_dir}")
+    log.info("=" * 60)
+
+
+if __name__ == "__main__":
+    main()