11 KiB

Raw Permalink Blame History

Round-Robin Training Across 6 K80 Dies

Concept

Train up to 6 models/adapters concurrently by time-slicing across all 6 K80 dies. Each die holds its own model + optimizer state in VRAM persistently. A scheduler daemon cycles through dies, running one training step per visit, then moving on.

Not truly parallel (one die computes at a time) but concurrent — all 6 make progress. The 0.3ms swap cost is negligible compared to training step time.

Why It Works

VRAM persists across bar_swap release/acquire (Architecture A, confirmed)
CUDA context survives the swap — no weight re-upload, no driver reload
Each die has 12GB VRAM (K80) — enough for models up to ~6B INT8 with optimizer state
Swap overhead: 0.3ms per die switch vs hundreds of ms per training step

Architecture

┌─────────────────────────────────────────────┐
│              Scheduler Daemon                │
│                                              │
│  for each round:                             │
│    for die_id in 0..5:                       │
│      1. bar_swap acquire(die_id)    ~0.3ms   │
│      2. set CUDA context for die             │
│      3. run N training steps                 │
│      4. bar_swap release(die_id)    ~0.1ms   │
│                                              │
│  repeat until all jobs report done           │
└─────────────────────────────────────────────┘

VRAM per die (persistent):
  ┌──────────────┐  ┌──────────────┐  ┌──────────────┐
  │   Die 0      │  │   Die 1      │  │   Die 2      │
  │  3.6B LoRA   │  │  3.6B LoRA   │  │  0.2B full   │
  │  adapter A   │  │  adapter B   │  │  finetune    │
  │  + optimizer │  │  + optimizer │  │  + optimizer │
  └──────────────┘  └──────────────┘  └──────────────┘
  ┌──────────────┐  ┌──────────────┐  ┌──────────────┐
  │   Die 3      │  │   Die 4      │  │   Die 5      │
  │  3.6B LoRA   │  │  1.5B RFA    │  │  0.2B LoRA   │
  │  adapter C   │  │  hybrid      │  │  adapter D   │
  │  + optimizer │  │  + optimizer │  │  + optimizer │
  └──────────────┘  └──────────────┘  └──────────────┘

VRAM Budget Per Die (12GB)

Model Size	Weights (INT8)	Optimizer State (LoRA r=32)	Activations	Fits?
0.2B	~200 MB	~50 MB	~100 MB	Easily (11.6 GB free)
1.5B	~1.5 GB	~200 MB	~500 MB	Yes (9.8 GB free)
3.6B	~3.6 GB	~400 MB	~1 GB	Yes (7 GB free)
7B	~7 GB	~800 MB	~2 GB	Tight (~2 GB free)

Training Step Timing Estimates

Per-step time depends on model size, sequence length, and batch size. These are estimates for single-example training steps on K80 (sm_37, ~240 GB/s memory bandwidth, ~2.8 TFLOPS FP32).

Model	Seq Len	Est. Step Time	Steps/min (1 die)	Steps/min (6 dies round-robin)
0.2B	256	~200ms	~300	~270 (90% efficient)
0.2B	512	~400ms	~150	~140
3.6B	256	~2s	~30	~29.5 (98% efficient)
3.6B	512	~4s	~15	~14.9

Round-robin efficiency = step_time / (step_time + swap_overhead). With 0.3ms swap, even the fastest step (200ms) is 99.85% efficient per die. The main cost is that each die only gets 1/6 of the wall-clock time.

Steps Per Visit (Amortization)

Running more than 1 step per die visit amortizes swap cost further and improves GPU utilization. Trade-off: longer time before other dies get their turn.

Steps/Visit	Round Time (6 dies, 3.6B)	Swap Overhead
1	~12s	0.015%
10	~120s	0.0015%
100	~20min	negligible

For overnight runs, 100+ steps per visit makes sense. For interactive monitoring, 1-10 steps per visit gives faster feedback across all jobs.

Scheduler Design

Job Config

typedef struct {
    int die_id;                     // which die (0-5)
    const char *model_path;         // base model weights
    const char *data_path;          // training data
    const char *output_path;        // where to save adapter
    int lora_rank;                  // LoRA rank
    float lr;                       // learning rate
    int max_epochs;                 // when to stop
    int max_seq_len;                // sequence length
    int steps_per_visit;            // steps per round-robin slot
    bool quantize;                  // INT8 base weights
} rr_job;

Daemon Lifecycle

1. INIT PHASE (one-time, sequential):
   for each job:
     a. acquire die
     b. load model + create CUDA context + upload weights
     c. create LoRA adapter + optimizer
     d. release die

   All 6 dies now hold their model + optimizer in VRAM.

2. TRAINING PHASE (round-robin loop):
   while any job is not done:
     for each active job:
       a. acquire die                         ~0.3ms
       b. restore CUDA context
       c. run steps_per_visit training steps
       d. log loss, step count, throughput
       e. release die                         ~0.1ms
       f. if job.step >= job.max_steps: mark done

3. SAVE PHASE (sequential):
   for each completed job:
     a. acquire die
     b. download adapter weights from VRAM
     c. save to output_path as .safetensors
     d. release die

4. CLEANUP:
   destroy all CUDA contexts, unload nvidia

CLI Interface

# Config file: jobs.json
[
  {
    "die": 0,
    "model": "~/models/base/RWKV-X-3.6B-64k-Base.safetensors",
    "data": "~/models/datasets/code_instruct.jsonl",
    "output": "~/models/loras/code-3.6B/",
    "lora_rank": 32,
    "lr": 2e-5,
    "epochs": 3,
    "seq_len": 512,
    "steps_per_visit": 50,
    "quantize": true
  },
  {
    "die": 1,
    "model": "~/models/base/RWKV-X-3.6B-64k-Base.safetensors",
    "data": "~/models/datasets/chat_instruct.jsonl",
    "output": "~/models/loras/chat-3.6B/",
    "lora_rank": 32,
    "lr": 2e-5,
    "epochs": 3,
    "seq_len": 512,
    "steps_per_visit": 50,
    "quantize": true
  }
]

# Launch
sudo ./rr_train --config jobs.json

# Output (live):
[round 1] die 0: step 50/1500, loss=2.34, 28.5 steps/min
[round 1] die 1: step 50/1500, loss=2.41, 27.8 steps/min
[round 1] die 2: step 50/900,  loss=1.89, 145 steps/min
...
[round 30] die 0: DONE (1500 steps, final loss=0.82)
[round 30] die 1: step 1450/1500, loss=0.79
...
Saving die 0 adapter to ~/models/loras/code-3.6B/adapter.safetensors
All jobs complete.

Current State (What Exists vs What's Needed)

The training pipeline today splits across CPU and GPU:

Step	Where it runs now	What it does
Forward pass	GPU (CUDA)	Matmuls, activations, norms — already ported, 38 tok/s on K80
Backward pass	CPU only	Computes gradients via chain rule — same matmuls/element-wise ops as forward, just not ported yet
Optimizer step	CPU only	AdamW: updates momentum, variance, applies weight update — element-wise tensor math
Weight update	CPU only	Applies gradients to LoRA parameters — scaling and addition

The backward pass, optimizer, and weight updates are all standard tensor math — matmuls (fc_sgemm), element-wise ops (fc_add, fc_mul), scaling. These are the same operations we already have CUDA kernels for from the forward pass. Porting them is not new math, it's calling the GPU versions instead of the CPU versions.

Why it's CPU-only today: Training was built on Mac (M4) where Metal handles the forward pass and CPU handles the rest. Apple Silicon's CPU is fast enough that this wasn't a bottleneck. On the K80 server the CPU is much slower, so GPU training actually matters here.

What needs to happen: Port the backward pass in rx_train.c / rfa_train.c to use fc_sgemm, fc_add, fc_mul etc. instead of tm_matmul, tm_add, tm_mul. Then move the AdamW optimizer step to GPU (element-wise update of momentum/variance/weights). Once both are on GPU, the entire training loop stays in VRAM and round-robin scheduling works.

Implementation Order

Phase 1: Single-die GPU training

Port backward pass to CUDA (use existing fc_* kernels — same math, different dispatch)
Port AdamW optimizer step to CUDA (element-wise: m = beta1*m + (1-beta1)*grad, etc.)
Keep optimizer state (m, v tensors) in VRAM via fc_buffer
This is the critical prerequisite — everything else depends on it

Phase 2: Multi-die state persistence

Prove that optimizer state (Adam momentum/variance) survives hot swap
Prove that gradient accumulators survive hot swap
Test: load model on die 0, train 10 steps, swap to die 1, swap back, train 10 more, compare loss curve to uninterrupted run

Phase 3: Round-robin scheduler

Implement the daemon loop
Job config parsing (JSON or simple text format)
Per-die CUDA context management
Logging and progress reporting

Phase 4: Robustness

Handle die acquire failures (skip die, retry next round)
Checkpoint adapters periodically (not just at end)
Resume from checkpoint after crash/reboot
Graceful shutdown (SIGINT saves all adapters)

Open Questions

GPU backward pass: Current training backward pass is CPU-only. Need to port gradient computation to CUDA kernels before any of this works. This is the biggest prerequisite.
Optimizer state in VRAM: AdamW needs momentum (m) and variance (v) tensors per parameter. For LoRA r=32 on 3.6B, this is ~800MB. Fits in 12GB but needs CUDA-side optimizer step.
Multiple CUDA contexts: Can we have 6 CUDA contexts (one per die) alive simultaneously? Current code creates one context in fc_init(). May need to refactor to support context-per-die.
Gradient checkpointing on GPU: Current training uses CPU gradient checkpointing. GPU training will need GPU-side checkpointing to fit in 12GB for larger models.
P40 compatibility: P40 is single-die (24GB), Pascal (sm_61). bar_swap would treat each card as 1 die. Same scheduler works, just fewer dies per card but more VRAM each. CUDA 11.4 supports sm_61. PTX kernels need recompile for sm_61 (or use sm_37 PTX which JIT-compiles up).

11 KiB Raw Permalink Blame History