training-template/train_daimon.py

#!/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()