Ai Model Training

Model Training Agent

Purpose

Design and execute model training plans for LLM fine-tuning, continued pre-training, and RLHF alignment: strategy selection, data pipeline, training configuration, distributed setup, hyperparameter optimization, evaluation, and production tracking.

Agent Protocol

Trigger

User request includes: fine-tuning, LoRA, QLoRA, RLHF, DPO, PPO, training LLM, model training, instruction tuning, preference tuning, SFT, prompt tuning, adapter, PEFT, Supervised Fine-Tuning, distributed training, hyperparameter search, pre-training, continued pre-training.

Protocol

1. Clarify: base model, task type, data volume (size + tokens), compute budget (GPU hours, dollars), hardware available.
2. Navigate decision tree to select training approach.
3. Prepare training data: format (instruction / chat / preference pairs), tokenize, split, validate.
4. Configure training: hyperparameters, optimizer, LR schedule, precision, batch size.
5. Design distributed setup: single GPU, FSDP, DeepSpeed, multi-node.
6. Define evaluation: pre-training baseline, in-training metrics, post-training benchmarks, forgetting checks.
7. Set up experiment tracking: metrics logging, checkpoint registry, hyperparameter capture.

Decision Tree: Training Approach

Q: Is this your first time training this model?
├── NO  → Go to "Fine-tuning or continued pre-training?"
└── YES → Go to "Available compute?"

Q: Available compute?
├──  160 GB VRAM / multi-node → Full fine-tune (distributed)
    BF16, FSDP or DeepSpeed ZeRO-3, tensor parallelism for 70B+

Q: Fine-tuning or continued pre-training?
├── Task adaptation ( 100K examples) → Full fine-tune
├── New knowledge / continued pre-training       → Full pre-train or continued pre-train
└── Align model behavior                         → Go to "Alignment method?"

Q: Alignment method?
├── Human preference data available?
│   ├── YES → Ask: KL control importance?
│   │   ├── HIGH → PPO (3-stage: SFT → RM → PPO)
│   │   └── LOW  → DPO (single stage, no reward model)
│   └── NO  → SFT only (instruction tuning)
└── Want to avoid reward model training?
    ├── YES → DPO
    └── NO  → PPO (if compute budget allows 3 stages)

Q: Multi-task / multi-domain?
├── YES → Use LoRA adapters per task with shared base
│   Consider: AdapterFusion, LoRA ensembles
└── NO  → Single adapter or full fine-tune

Q: Data size for instruction tuning?
├──  100K examples → Full fine-tune preferred, 1-2 epochs

Workflow

Step 1: Select Training Method

• Full Pre-training: Train from scratch. Requires massive data (1T+ tokens), compute, and engineering. Only when no suitable base model exists.
• Continued Pre-training: Train on existing base with new domain data (code, biomedical, legal). Use same tokenizer, extend vocab if needed. LR 1e-5 to 5e-5.
• Full Fine-tune: All parameters updated. Best for large distribution shifts. Requires most compute. LR 1e-5 to 5e-5.
• LoRA: Low-rank adapters. ~1% of parameters. Best for task adaptation. Default choice. LR 1e-4 to 5e-4.
• QLoRA: Quantized LoRA (4-bit NF4) with double quantization. ~0.5% of parameters. Best for limited GPU memory. LR 1e-4 to 3e-4.
• Adapters: Bottleneck layers between transformer sublayers. Best for multi-task setups.
• DPO: Direct Preference Optimization. Single-stage alignment. No reward model needed.
• PPO: 3-stage RLHF. SFT → Reward Model → PPO. Most compute but best alignment control.

Step 2: Prepare Training Data

Instruction Format

from datasets import Dataset

data = [
    {"instruction": "Translate to French", "input": "Hello world", "output": "Bonjour le monde"},
    {"instruction": "Summarize", "input": "Long text...", "output": "Short summary..."}
]

Chat Template Format

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct")
tokenizer.pad_token = tokenizer.eos_token
messages = [
    {"role": "user", "content": "Hello"},
    {"role": "assistant", "content": "Hi there!"}
]
formatted = tokenizer.apply_chat_template(messages, tokenize=False)

Preference Pairs (for DPO/RLHF)

preference_data = [
    {
        "prompt": "What is the capital of France?",
        "chosen": "Paris is the capital of France.",
        "rejected": "London is the capital of France."
    }
]

Tokenization with Label Masking

def tokenize_and_mask(examples, tokenizer, max_length=2048):
    outputs = tokenizer(
        examples["text"],
        truncation=True,
        max_length=max_length,
        padding="max_length",
        return_tensors=None,
    )
    # Copy input_ids to labels, mask user tokens with -100
    outputs["labels"] = outputs["input_ids"].copy()
    return outputs

Data Splitting & Validation

# Split: train (80%), eval (10%), test (10%)
# Stratify by category if available
from sklearn.model_selection import train_test_split

def prepare_splits(data, stratify_col=None):
    train_val, test = train_test_split(
        data, test_size=0.1, stratify=stratify_col
    )
    train, eval = train_test_split(
        train_val, test_size=0.111, stratify=stratify_col
    )
    return train, eval, test

Step 3: Configure Training with LoRA

from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, TrainingArguments, Trainer
import torch

model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

Step 4: Training Arguments

training_args = TrainingArguments(
    output_dir="./checkpoints",
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=8,
    learning_rate=2e-4,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    num_train_epochs=3,
    logging_steps=10,
    logging_strategy="steps",
    evaluation_strategy="steps",
    eval_steps=200,
    save_strategy="steps",
    save_steps=500,
    save_total_limit=3,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    bf16=True,
    tf32=True,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    optim="adamw_torch",
    weight_decay=0.01,
    max_grad_norm=1.0,
    report_to="wandb",
    run_name=f"lora-ft-{model_name}-{timestamp}",
    remove_unused_columns=False,
    dataloader_num_workers=4,
    ddp_find_unused_parameters=False,
)

Step 5: Trainer Setup

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
    data_collator=lambda data: tokenizer.pad(
        [{"input_ids": d["input_ids"], "attention_mask": d["attention_mask"], "labels": d["labels"]} for d in data],
        return_tensors="pt",
    ),
    compute_metrics=compute_metrics_fn if eval_task else None,
)
trainer.train()

Step 6: Distributed Training

FSDP Configuration

# fsdp_config.yaml
compute_environment: LOCAL_MACHINE
distributed_type: FSDP
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
  fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

DeepSpeed ZeRO-3 Configuration

{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {"device": "cpu", "pin_memory": true},
        "offload_param": {"device": "cpu", "pin_memory": true},
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": 5e7,
        "stage3_prefetch_bucket_limit": 5e7,
        "stage3_param_persistence_threshold": 1e6,
        "sub_group_size": 1e9
    },
    "bf16": {"enabled": true},
    "fp16": {"enabled": false},
    "gradient_accumulation_steps": 8,
    "gradient_clipping": 1.0,
    "steps_per_print": 100,
    "train_batch_size": 32,
    "train_micro_batch_size_per_gpu": 4,
    "wall_clock_breakdown": false
}

Launch Commands

# DeepSpeed
deepspeed --num_gpus=8 train.py \
    --deepspeed ds_config.json \
    --model_name meta-llama/Llama-2-13b-hf

# FSDP via torchrun
torchrun --nproc_per_node=8 train.py \
    --fsdp full_shard \
    --fsdp_transformer_layer_cls_to_wrap LlamaDecoderLayer

# Multi-node
torchrun --nnodes=4 --nproc_per_node=8 --rdzv_id=101 --rdzv_backend=c10d train.py

Architectural Patterns

Data Pipeline Architecture

Raw Sources (JSONL, Parquet, DB)
  → Data Cleaner (PII removal, dedup, quality scoring)
    → Formatter (chat template, instruction format)
      → Tokenizer (map-style dataset with caching)
        → DataLoader (batching, shuffling, num_workers)
          → Training Loop

Key design decisions:

• Use datasets library with memory mapping for large datasets (no full RAM load).
• Cache tokenized datasets to disk (keep_in_memory=False) between runs.
• Set dataloader_num_workers=4-8 to avoid GPU starvation.
• Use StreamingDataset for datasets larger than available disk.

Training Loop Architecture

# Custom training loop (when Trainer is insufficient)
for epoch in range(num_epochs):
    for step, batch in enumerate(train_dataloader):
        batch = {k: v.to(device) for k, v in batch.items()}

        with ctx:  # autocast for mixed precision
            outputs = model(**batch)
            loss = outputs.loss / gradient_accumulation_steps

        loss_scaler.scale(loss).backward()

        if (step + 1) % gradient_accumulation_steps == 0:
            loss_scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            loss_scaler.step(optimizer)
            loss_scaler.update()
            optimizer.zero_grad()

        if step % logging_steps == 0:
            metrics = compute_metrics(model, eval_loader, device)
            log_to_tracker({"train/loss": loss.item(), **metrics})

Checkpointing Architecture

class CheckpointManager:
    def __init__(self, output_dir, save_every_n_steps, keep_last_k):
        self.dir = output_dir
        self.save_every = save_every_n_steps
        self.keep = keep_last_k
        self.checkpoints = []

    def save(self, model, optimizer, scheduler, step, metrics):
        if step % self.save_every != 0 and not self._is_best(metrics):
            return
        ckpt_path = os.path.join(self.dir, f"step_{step}")
        os.makedirs(ckpt_path, exist_ok=True)
        save_args = {
            "state_dict": model.state_dict(),
            "optimizer": optimizer.state_dict(),
            "scheduler": scheduler.state_dict(),
            "step": step,
            "metrics": metrics,
        }
        torch.save(save_args, os.path.join(ckpt_path, "training_state.pt"))
        model.save_pretrained(ckpt_path)
        self.checkpoints.append((step, metrics.get("eval_loss", float("inf"))))
        self.checkpoints.sort(key=lambda x: x[1])
        while len(self.checkpoints) > self.keep:
            stale_step = self.checkpoints.pop()[0]
            shutil.rmtree(os.path.join(self.dir, f"step_{stale_step}"))

Evaluation Loop Architecture

# In-training evaluation
# Run on a fixed subset of eval data (500-1000 samples) every N steps
# Track: eval_loss, perplexity, task accuracy, gradient norms
#
# Pre-training baseline: run before training starts
# In-training: every N steps on eval subset
# Post-training: full benchmark suite after training
#
# Catastrophic forgetting detection:
# - Maintain a "forgetting set" of diverse tasks
# - Track scores relative to pre-training baseline
# - Alert if any task drops > 5% from baseline

def evaluate_model(model, eval_dataset, tokenizer, device, max_samples=500):
    model.eval()
    total_loss = 0.0
    total_steps = 0
    with torch.no_grad():
        for batch in islice(eval_dataloader, max_samples // batch_size):
            batch = {k: v.to(device) for k, v in batch.items()}
            outputs = model(**batch)
            total_loss += outputs.loss.item()
            total_steps += 1
    return {"eval_loss": total_loss / total_steps, "perplexity": math.exp(total_loss / total_steps)}

Training Infrastructure Design

Compute

| GPU | VRAM | FP16 TFLOPS | BF16 TFLOPS | Best For | |-----|------|-------------|-------------|----------| | RTX 4090 | 24 GB | 82 | N/A | QLoRA, LoRA ≤13B | | A100 80GB | 80 GB | 312 | 312 | Full FT ≤13B, LoRA ≤70B | | H100 | 80 GB | 989 | 989 | Full FT ≤70B, pre-training | | H200 | 141 GB | 989 | 989 | Full FT ≤70B+ | | MI300X | 192 GB | 653 | 653 | Alternative to H100 |

Storage Requirements

• Dataset storage: NVMe SSD recommended. Tokenized datasets benefit from fast random access.
• Checkpoint storage: Large contiguous writes. One 70B checkpoint = ~140 GB (BF16) or ~560 GB (optimizer states + model). Budget 3-5x model size for checkpoint space.
• Model registry: Object storage (S3, GCS, Blob) for versioned artifacts.
• Cache: HuggingFace cache directory needs 10-100 GB for base models.

Networking (Multi-Node)

• Minimum: 25 Gbps Ethernet. Expect 30-40% scaling efficiency.
• Recommended: 200-400 Gbps InfiniBand (HDR/HDR100/NDR). Expect 80-90% scaling efficiency.
• Topology: Fat-tree or Dragonfly for GPU clusters.
• NCCL: Use NCCL_IB_HCA, NCCL_SOCKET_IFNAME, tune NCCL_IB_TIMEOUT and NCCL_IB_RETRY_CNT.

CPU/RAM Guidelines

• Per GPU: Minimum 64 GB system RAM per GPU (128 GB recommended for ZeRO-3 with CPU offload).
• CPU cores: At least 8-16 cores per GPU for data loading and preprocessing.

Hyperparameter Optimization Strategies

Bayesian Optimization with Optuna

import optuna
from optuna.integration import TransformersPruningCallback

def objective(trial):
    lr = trial.suggest_float("learning_rate", 5e-5, 5e-4, log=True)
    lora_r = trial.suggest_int("lora_r", 8, 64)
    lora_alpha = trial.suggest_int("lora_alpha", 16, 128)
    weight_decay = trial.suggest_float("weight_decay", 0.0, 0.1)
    warmup_ratio = trial.suggest_float("warmup_ratio", 0.01, 0.1)
    dropout = trial.suggest_float("lora_dropout", 0.0, 0.3)

    config = LoraConfig(r=lora_r, lora_alpha=lora_alpha, lora_dropout=dropout)
    model = get_peft_model(base_model, config)

    args = TrainingArguments(
        learning_rate=lr,
        weight_decay=weight_decay,
        warmup_ratio=warmup_ratio,
        num_train_epochs=2,
        report_to="none",
        logging_steps=50,
        save_strategy="no",
    )
    trainer = Trainer(model=model, args=args, train_dataset=train_data, eval_dataset=eval_data)
    trainer.train()
    eval_result = trainer.evaluate()
    return eval_result["eval_loss"]

study = optuna.create_study(direction="minimize", pruner=optuna.pruners.MedianPruner())
study.optimize(objective, n_trials=20)
print(f"Best params: {study.best_params}, best loss: {study.best_value}")

Learning Rate Range Test

# Find optimal LR by running a short training with increasing LR
# Use lr_finder from transformers or implement manually
def lr_range_test(model, dataloader, optimizer_cls, device, min_lr=1e-7, max_lr=1):
    optimizer = optimizer_cls(model.parameters(), lr=min_l

…

## Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

- **Author:** [j4flmao](https://github.com/j4flmao)
- **Source:** [j4flmao/agent-skills](https://github.com/j4flmao/agent-skills)
- **License:** MIT

Install and usage instructions live in the source repository linked above.