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