Add training template: test_template.py

training-template/test_template.py

@@ -0,0 +1,549 @@
+#!/usr/bin/env python3
+"""
+Daimon Training Template — Validation Tests
+=============================================
+
+Run these BEFORE launching training to catch configuration issues early.
+Each test is independent and reports PASS/FAIL.
+
+Usage:
+    python3 /workspace/runpod-template/test_template.py
+
+These tests run on a SINGLE GPU (no distributed launch needed).
+"""
+
+import os
+import sys
+import json
+import time
+import traceback
+
+RESULTS = []
+
+
+def test(name):
+    """Decorator to register and run a test."""
+    def decorator(fn):
+        def wrapper():
+            try:
+                fn()
+                RESULTS.append(("PASS", name, None))
+                print(f"  PASS: {name}")
+            except Exception as e:
+                RESULTS.append(("FAIL", name, str(e)))
+                print(f"  FAIL: {name}")
+                print(f"        {e}")
+                traceback.print_exc()
+        wrapper.__name__ = name
+        wrapper._test = True
+        return wrapper
+    return decorator
+
+
+# ── Test 1: GPU with sufficient VRAM ──────────────────────────────────────
+
+@test("GPU with >= 140GB VRAM is available")
+def test_gpu_vram():
+    import torch
+    gpu_count = torch.cuda.device_count()
+    assert gpu_count >= 1, (
+        f"No GPUs detected. Need at least 1x H200 SXM 141GB."
+    )
+
+    # Find the GPU with the most VRAM
+    max_vram_gb = 0
+    for i in range(gpu_count):
+        name = torch.cuda.get_device_name(i)
+        mem_gb = torch.cuda.get_device_properties(i).total_memory / 1e9
+        max_vram_gb = max(max_vram_gb, mem_gb)
+        print(f"        GPU {i}: {name}, {mem_gb:.1f} GB")
+
+    assert max_vram_gb >= 140, (
+        f"Largest GPU has {max_vram_gb:.1f} GB VRAM. Need >= 140 GB (H200 SXM). "
+        f"Full SFT needs ~90GB GPU (70GB model + 20GB activations)."
+    )
+
+
+# ── Test 2: System RAM >= 180GB (critical for CPU offload) ───────────────
+
+@test("System RAM >= 180GB for CPU-offloaded optimizer")
+def test_system_ram():
+    """
+    Full-parameter SFT offloads gradients (~70GB) and Adafactor states (~35GB)
+    to CPU RAM. Without enough system RAM, training will OOM on the CPU side.
+    """
+    ram_bytes = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
+    ram_gb = ram_bytes / 1e9
+    print(f"        System RAM: {ram_gb:.0f} GB")
+
+    assert ram_gb >= 180, (
+        f"System RAM is {ram_gb:.0f} GB. Need >= 180 GB. "
+        f"CPU-offloaded memory budget: gradients (~70GB) + Adafactor (~35GB) = ~105GB, "
+        f"plus OS and data loading overhead. "
+        f"AdamW would need ~280GB — that's why we use Adafactor."
+    )
+
+    # Warn if tight
+    if ram_gb < 200:
+        print(f"        WARNING: {ram_gb:.0f}GB is tight. 200GB+ recommended.")
+        print(f"        CPU budget: ~105GB for offloaded states + ~30GB overhead")
+
+
+# ── Test 3: Full SFT config is valid (no LoRA) ──────────────────────────
+
+@test("Full SFT config is valid (no LoRA)")
+def test_sft_config():
+    import yaml
+
+    config_paths = [
+        "/workspace/runpod-template/train_daimon_config.yaml",
+        os.path.join(os.path.dirname(__file__), "train_daimon_config.yaml"),
+    ]
+    found = None
+    for p in config_paths:
+        if os.path.exists(p):
+            found = p
+            break
+
+    assert found is not None, (
+        f"Config not found. Looked in: {config_paths}"
+    )
+
+    with open(found) as f:
+        config = yaml.safe_load(f)
+
+    # Verify NO LoRA section — this is full-parameter SFT
+    assert "lora" not in config, (
+        "Config still has a 'lora' section. This template does full-parameter SFT — "
+        "remove the lora section entirely."
+    )
+
+    # Verify Adafactor optimizer is configured
+    optimizer = config.get("optimizer", "")
+    assert optimizer.lower() == "adafactor", (
+        f"Optimizer must be 'adafactor' for full SFT on single node. "
+        f"Got: '{optimizer}'. AdamW needs ~280GB CPU RAM — not viable."
+    )
+
+    # Verify DeepSpeed config is referenced
+    ds_config = config.get("deepspeed_config")
+    assert ds_config is not None, (
+        "Config must reference deepspeed_config for ZeRO-2 CPU offload. "
+        "Full-parameter SFT cannot fit without gradient sharding."
+    )
+
+    # Verify learning rate is appropriate for full SFT (not LoRA rate)
+    lr = config.get("learning_rate", 0)
+    assert lr <= 1e-4, (
+        f"Learning rate {lr} is too high for full SFT. "
+        f"LoRA uses 2e-4, but full SFT should be 1e-5 to 5e-6 to avoid "
+        f"destabilizing MoE routing gates."
+    )
+
+    # Verify model revision is pinned
+    revision = config.get("model_revision")
+    assert revision is not None and len(revision) >= 10, (
+        f"model_revision should be pinned to a specific commit hash. Got: {revision}"
+    )
+
+    # Verify bf16 is enabled
+    assert config.get("bf16") is True, "bf16 must be enabled"
+
+    # Verify save_total_limit is small (full checkpoints are ~70GB each)
+    save_limit = config.get("save_total_limit", 10)
+    assert save_limit <= 5, (
+        f"save_total_limit={save_limit} is too high for full SFT. "
+        f"Each checkpoint is ~70GB. Limit to 3-5 to avoid filling disk."
+    )
+
+    print(f"        Config: {found}")
+    print(f"        Optimizer: {optimizer}")
+    print(f"        DeepSpeed: {ds_config}")
+    print(f"        Learning rate: {lr}")
+    print(f"        Model revision: {revision[:12]}...")
+    print(f"        bf16: {config.get('bf16')}")
+    print(f"        save_total_limit: {save_limit}")
+    print(f"        Method: Full-parameter SFT (no LoRA)")
+
+
+# ── Test 4: DeepSpeed ZeRO-2 config exists and is valid ──────────────────
+
+@test("DeepSpeed ZeRO-2 config is valid")
+def test_deepspeed_config():
+    ds_paths = [
+        "/workspace/runpod-template/ds_config_zero2.json",
+        os.path.join(os.path.dirname(__file__), "ds_config_zero2.json"),
+    ]
+    found = None
+    for p in ds_paths:
+        if os.path.exists(p):
+            found = p
+            break
+
+    assert found is not None, (
+        f"DeepSpeed config not found. Looked in: {ds_paths}"
+    )
+
+    with open(found) as f:
+        ds_config = json.load(f)
+
+    # Verify it's ZeRO Stage 2 (not 3)
+    stage = ds_config.get("zero_optimization", {}).get("stage")
+    assert stage == 2, (
+        f"DeepSpeed must be ZeRO Stage 2, got stage {stage}. "
+        f"Stage 2 shards gradients; Stage 3 shards params too (not needed for single GPU "
+        f"where the model fits in VRAM)."
+    )
+
+    # Verify optimizer offload to CPU
+    offload_opt = ds_config.get("zero_optimization", {}).get("offload_optimizer", {})
+    assert offload_opt.get("device") == "cpu", (
+        f"Optimizer must be offloaded to CPU. "
+        f"Adafactor states (~35GB) need to live in system RAM."
+    )
+
+    # Verify params stay on GPU (not offloaded)
+    offload_param = ds_config.get("zero_optimization", {}).get("offload_param", {})
+    param_device = offload_param.get("device", "none")
+    assert param_device == "none", (
+        f"Parameters should NOT be offloaded (device={param_device}). "
+        f"The model fits in GPU VRAM — offloading params to CPU would be slow."
+    )
+
+    # Verify NO optimizer configured in DeepSpeed (we handle Adafactor in the script)
+    assert "optimizer" not in ds_config, (
+        "DeepSpeed config should NOT have an optimizer section. "
+        "Adafactor is configured in the training script directly — "
+        "DeepSpeed's optimizer config conflicts with custom optimizers."
+    )
+
+    # Verify bf16 enabled
+    assert ds_config.get("bf16", {}).get("enabled") is True, "bf16 must be enabled in DeepSpeed config"
+
+    print(f"        Config: {found}")
+    print(f"        ZeRO Stage: {stage}")
+    print(f"        Optimizer offload: CPU (pin_memory={offload_opt.get('pin_memory')})")
+    print(f"        Param offload: none (stays on GPU)")
+    print(f"        bf16: enabled")
+    print(f"        No DeepSpeed optimizer block (Adafactor managed by script)")
+
+
+# ── Test 5: Model architecture loads ─────────────────────────────────────
+
+@test("Model architecture loads")
+def test_model_loads():
+    import yaml
+    from transformers import AutoConfig, AutoTokenizer
+
+    # Read revision from config
+    config_paths = [
+        "/workspace/runpod-template/train_daimon_config.yaml",
+        os.path.join(os.path.dirname(__file__), "train_daimon_config.yaml"),
+    ]
+    revision = None
+    for p in config_paths:
+        if os.path.exists(p):
+            with open(p) as f:
+                cfg = yaml.safe_load(f)
+            revision = cfg.get("model_revision")
+            break
+
+    model_id = "Qwen/Qwen3.6-35B-A3B"
+    local_path = "/workspace/models/Qwen3.6-35B-A3B"
+
+    if os.path.isdir(local_path) and os.path.exists(f"{local_path}/config.json"):
+        source = local_path
+    else:
+        source = model_id
+
+    kwargs = {"trust_remote_code": True}
+    if revision and source == model_id:
+        kwargs["revision"] = revision
+
+    config = AutoConfig.from_pretrained(source, **kwargs)
+    print(f"        Model: {source}")
+    print(f"        Type: {config.model_type}")
+    print(f"        Hidden: {config.hidden_size}")
+    print(f"        Layers: {config.num_hidden_layers}")
+    print(f"        Experts: {getattr(config, 'num_experts', 'N/A')}")
+
+    # Also verify tokenizer loads
+    tokenizer = AutoTokenizer.from_pretrained(source, **kwargs)
+    print(f"        Vocab: {tokenizer.vocab_size}")
+
+
+# ── Test 6: Training data loads and sequences are within bounds ────────────
+
+@test("Training data loads with valid sequence lengths")
+def test_data_loads():
+    import yaml
+    from datasets import load_from_disk
+    from transformers import AutoTokenizer
+
+    data_dir = "/workspace/daimon-data"
+    train_arrow = f"{data_dir}/train_arrow"
+
+    assert os.path.isdir(train_arrow), (
+        f"Training data not found at {train_arrow}. "
+        f"Run setup.sh first to download and prepare data."
+    )
+
+    train_ds = load_from_disk(train_arrow)
+    print(f"        Train samples: {len(train_ds):,}")
+
+    # Check a sample
+    sample = train_ds[0]
+    assert "messages" in sample, f"Expected 'messages' key, got: {list(sample.keys())}"
+    assert len(sample["messages"]) >= 2, "Each sample needs at least 2 messages (user + assistant)"
+
+    # Verify sequence lengths against max_seq_length
+    model_id = "Qwen/Qwen3.6-35B-A3B"
+    local_path = "/workspace/models/Qwen3.6-35B-A3B"
+    source = local_path if os.path.isdir(local_path) else model_id
+
+    # Read revision from config
+    config_paths = [
+        "/workspace/runpod-template/train_daimon_config.yaml",
+        os.path.join(os.path.dirname(__file__), "train_daimon_config.yaml"),
+    ]
+    kwargs = {"trust_remote_code": True}
+    for p in config_paths:
+        if os.path.exists(p):
+            with open(p) as f:
+                cfg = yaml.safe_load(f)
+            revision = cfg.get("model_revision")
+            if revision and source == model_id:
+                kwargs["revision"] = revision
+            break
+
+    tokenizer = AutoTokenizer.from_pretrained(source, **kwargs)
+
+    max_seq_length = 4096  # From config (reduced for full SFT)
+    too_long = 0
+    max_found = 0
+
+    for i, example in enumerate(train_ds):
+        try:
+            text = tokenizer.apply_chat_template(
+                example["messages"], tokenize=False, add_generation_prompt=False
+            )
+            tokens = len(tokenizer.encode(text, add_special_tokens=False))
+            max_found = max(max_found, tokens)
+            if tokens > max_seq_length:
+                too_long += 1
+        except Exception:
+            pass
+
+        if i >= 100:  # Check first 100 samples
+            break
+
+    print(f"        Max tokens in sample: {max_found}")
+    print(f"        Exceeding {max_seq_length}: {too_long}/{min(len(train_ds), 101)}")
+
+    if too_long > 0:
+        print(f"        WARNING: {too_long} sequences exceed max_seq_length.")
+        print(f"        The training script will pre-split these, but check your data.")
+
+
+# ── Test 7: Persistent volume is mounted and writable ──────────────────────
+
+@test("Persistent volume is mounted and writable")
+def test_persistent_volume():
+    workspace = "/workspace"
+    assert os.path.isdir(workspace), "/workspace not found. Is the persistent volume mounted?"
+
+    # Check it's writable
+    test_file = os.path.join(workspace, ".daimon_write_test")
+    try:
+        with open(test_file, "w") as f:
+            f.write("test")
+        os.remove(test_file)
+    except PermissionError:
+        raise AssertionError("/workspace is not writable. Check volume permissions.")
+
+    # Check available space — full checkpoints are ~70GB each
+    import shutil
+    total, used, free = shutil.disk_usage(workspace)
+    free_gb = free / (1024**3)
+    total_gb = total / (1024**3)
+    print(f"        Volume: {total_gb:.0f} GB total, {free_gb:.0f} GB free")
+
+    assert free_gb >= 200, (
+        f"Only {free_gb:.0f} GB free on /workspace. "
+        f"Need at least 200GB for model + full checkpoints (~70GB each, limit=3)."
+    )
+
+
+# ── Test 8: Memory estimate for full SFT ────────────────────────────────
+
+@test("Memory estimate: full SFT fits in GPU + CPU")
+def test_memory_estimate():
+    """
+    Estimate memory usage for full-parameter SFT with Adafactor + ZeRO-2.
+    Verifies both GPU VRAM and system RAM are sufficient.
+    Does NOT load the full model — just calculates from config.
+    """
+    import torch
+
+    # Qwen3.6-35B-A3B has ~35B total params
+    total_params = 35e9
+
+    # GPU memory budget
+    model_gb = total_params * 2 / 1e9         # bf16 = 2 bytes per param = ~70GB
+    activation_gb = 20.0                        # with gradient checkpointing
+    gpu_total = model_gb + activation_gb        # ~90GB
+
+    # CPU memory budget (ZeRO-2 offloaded)
+    gradient_gb = total_params * 2 / 1e9       # bf16 gradients = ~70GB
+    # Adafactor: factored second moments, roughly 1 state per param in mixed precision
+    # Much less than AdamW's 2 fp32 states (280GB)
+    adafactor_gb = total_params * 1 / 1e9      # ~35GB (conservative estimate)
+    cpu_total = gradient_gb + adafactor_gb      # ~105GB
+
+    # AdamW comparison (for reference)
+    adamw_gb = total_params * 4 * 2 / 1e9      # 2 fp32 states = ~280GB
+
+    # Available resources
+    gpu_vram_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+    ram_bytes = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
+    system_ram_gb = ram_bytes / 1e9
+
+    print(f"        === GPU Memory ===")
+    print(f"        Model params (bf16):     {model_gb:.1f} GB")
+    print(f"        Activations (grad ckpt): {activation_gb:.1f} GB")
+    print(f"        GPU total:               {gpu_total:.1f} GB")
+    print(f"        GPU available:           {gpu_vram_gb:.1f} GB")
+    print(f"        GPU headroom:            {gpu_vram_gb - gpu_total:.1f} GB")
+    print(f"        ")
+    print(f"        === CPU Memory (offloaded) ===")
+    print(f"        Gradients (bf16):        {gradient_gb:.1f} GB")
+    print(f"        Adafactor states:        {adafactor_gb:.1f} GB")
+    print(f"        CPU total:               {cpu_total:.1f} GB")
+    print(f"        System RAM:              {system_ram_gb:.0f} GB")
+    print(f"        CPU headroom:            {system_ram_gb - cpu_total:.0f} GB")
+    print(f"        ")
+    print(f"        === Why NOT AdamW ===")
+    print(f"        AdamW states would need: {adamw_gb:.0f} GB CPU RAM")
+    print(f"        System RAM available:    {system_ram_gb:.0f} GB")
+    print(f"        Deficit:                 {adamw_gb - system_ram_gb:.0f} GB (does not fit)")
+
+    assert gpu_total < gpu_vram_gb, (
+        f"Estimated GPU usage ({gpu_total:.1f} GB) exceeds GPU capacity ({gpu_vram_gb:.1f} GB)."
+    )
+
+    assert cpu_total < system_ram_gb * 0.85, (
+        f"Estimated CPU usage ({cpu_total:.1f} GB) exceeds safe threshold "
+        f"({system_ram_gb * 0.85:.0f} GB = 85% of {system_ram_gb:.0f} GB). "
+        f"Need headroom for OS, data loading, and PyTorch buffers."
+    )
+
+
+# ── Test 9: Smoke test — imports and config validation ────────────────────
+
+@test("Training imports and config validation (smoke test)")
+def test_smoke():
+    """
+    Verify all training imports work and the config is valid.
+    Does NOT load the full model — that would require too much VRAM for a test.
+    """
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer, Adafactor
+
+    # Verify trainer imports (no peft needed)
+    print(f"        Verifying trainer imports...")
+    from trl import SFTTrainer, SFTConfig
+    import deepspeed
+
+    # Verify Adafactor is importable
+    print(f"        Adafactor: importable from transformers")
+
+    # Verify NO peft dependency
+    # (peft may be installed but should not be required)
+    print(f"        No peft/LoRA dependency required for full SFT")
+
+    # Create a minimal SFTConfig to verify all parameters are accepted
+    test_config = SFTConfig(
+        output_dir="/tmp/daimon_test",
+        max_length=256,
+        num_train_epochs=1,
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=1,
+        learning_rate=5e-6,
+        max_steps=1,
+        bf16=True,
+        gradient_checkpointing=True,
+        gradient_checkpointing_kwargs={"use_reentrant": False},
+        report_to="none",
+    )
+
+    print(f"        SFTConfig created successfully")
+    print(f"        DeepSpeed version: {deepspeed.__version__}")
+    print(f"        TRL version: {__import__('trl').__version__}")
+    print(f"        Transformers version: {__import__('transformers').__version__}")
+
+    # Verify YAML config loads
+    import yaml
+    config_paths = [
+        "/workspace/runpod-template/train_daimon_config.yaml",
+        os.path.join(os.path.dirname(__file__), "train_daimon_config.yaml"),
+    ]
+    for p in config_paths:
+        if os.path.exists(p):
+            with open(p) as f:
+                cfg = yaml.safe_load(f)
+            print(f"        YAML config loaded: {len(cfg)} keys")
+            break
+
+    # Clean up
+    import shutil
+    if os.path.exists("/tmp/daimon_test"):
+        shutil.rmtree("/tmp/daimon_test")
+
+    print(f"        Smoke test passed.")
+
+
+# ── Run all tests ──────────────────────────────────────────────────────────
+
+def main():
+    print("=" * 60)
+    print("  DAIMON TRAINING TEMPLATE — VALIDATION TESTS")
+    print("  Method: Full-Parameter SFT (no LoRA)")
+    print(f"  {time.strftime('%Y-%m-%dT%H:%M:%S')}")
+    print("=" * 60)
+    print()
+
+    # Collect all test functions
+    tests = [v for v in globals().values() if callable(v) and getattr(v, '_test', False)]
+
+    for test_fn in tests:
+        test_fn()
+        print()
+
+    # Summary
+    passed = sum(1 for r in RESULTS if r[0] == "PASS")
+    failed = sum(1 for r in RESULTS if r[0] == "FAIL")
+
+    print("=" * 60)
+    print(f"  RESULTS: {passed} passed, {failed} failed")
+    print()
+
+    for status, name, error in RESULTS:
+        marker = "PASS" if status == "PASS" else "FAIL"
+        print(f"    [{marker}] {name}")
+        if error:
+            print(f"           {error}")
+
+    print()
+    if failed == 0:
+        print("  STATUS: ALL TESTS PASSED — READY TO TRAIN")
+        print("  Next: bash /workspace/runpod-template/launch.sh")
+    else:
+        print("  STATUS: FIX FAILURES BEFORE TRAINING")
+
+    print("=" * 60)
+
+    sys.exit(failed)
+
+
+if __name__ == "__main__":
+    main()