@r1z4x/tezgah 1.0.0

package/skills/ai-research/01-model-architecture/litgpt/references/distributed-training.md

@@ -0,0 +1,452 @@
+<!-- vendored from orchestra-research/AI-research-SKILLs@773a52944ba4747a18bd4ae9ade53fff041adcbc 01-model-architecture/litgpt/references/distributed-training.md; MIT (c) Orchestra Research -->
+# Distributed Training
+
+Guide to FSDP (Fully Sharded Data Parallel) distributed training in LitGPT for scaling to multiple GPUs and nodes.
+
+## Overview
+
+LitGPT uses **Lightning Fabric** with **FSDP** to distribute training across multiple GPUs. FSDP shards model parameters, gradients, and optimizer states to enable training models larger than single-GPU memory.
+
+**When to use FSDP**:
+- Model doesn't fit on single GPU
+- Want faster training with multi-GPU
+- Training models >7B parameters
+- Need to scale across multiple nodes
+
+## Quick Start
+
+### Single Node Multi-GPU
+
+```bash
+# Train Llama 2 7B on 4 GPUs
+litgpt finetune_lora meta-llama/Llama-2-7b-hf \
+  --devices 4 \
+  --data JSON \
+  --data.json_path data/alpaca.json
+```
+
+FSDP is **automatically enabled** when `devices > 1`.
+
+### Multi-Node Training
+
+```bash
+# Train on 2 nodes with 8 GPUs each (16 total)
+litgpt finetune_lora meta-llama/Llama-2-70b-hf \
+  --devices 8 \
+  --num_nodes 2 \
+  --data JSON \
+  --data.json_path data/alpaca.json
+```
+
+## FSDP Configuration
+
+### Default FSDP Strategy
+
+When multiple devices are used, LitGPT applies this FSDP configuration:
+
+```python
+from lightning.fabric.strategies import FSDPStrategy
+from litgpt.model import Block
+
+strategy = FSDPStrategy(
+    auto_wrap_policy={Block},
+    state_dict_type="full",
+    sharding_strategy="HYBRID_SHARD"
+)
+```
+
+**Parameters**:
+- `auto_wrap_policy={Block}`: Automatically wraps each transformer `Block` with FSDP
+- `state_dict_type="full"`: Saves full model (assembled on rank 0) for easy deployment
+- `sharding_strategy="HYBRID_SHARD"`: Shards parameters, gradients, and optimizer states
+
+### Sharding Strategies
+
+| Strategy | Shards | Communication | Use Case |
+|----------|--------|---------------|----------|
+| `FULL_SHARD` (ZeRO-3) | Params + Grads + Optim | All-gather before forward/backward | Maximum memory savings |
+| `SHARD_GRAD_OP` (ZeRO-2) | Grads + Optim only | Reduce-scatter after backward | Faster than FULL_SHARD |
+| `HYBRID_SHARD` (default) | All (hybrid across nodes) | Optimized for multi-node | Best for clusters |
+| `NO_SHARD` | None | Broadcast | Single GPU (no FSDP) |
+
+**Recommendation**: Use default `HYBRID_SHARD` for multi-node, or `FULL_SHARD` for single-node multi-GPU.
+
+### State Dict Types
+
+| Type | Behavior | Use Case |
+|------|----------|----------|
+| `full` (default) | Gathers all shards on rank 0, saves single file | Easy deployment, inference |
+| `sharded` | Each rank saves its shard separately | Faster checkpointing, resume training |
+
+### Auto-Wrap Policy
+
+FSDP wraps model components based on `auto_wrap_policy`:
+
+```python
+auto_wrap_policy={Block}  # Wrap each transformer block
+```
+
+This means each `Block` (transformer layer) is independently sharded across GPUs. For a 32-layer model on 4 GPUs, each GPU holds ~8 layer shards.
+
+## Thunder FSDP (Advanced)
+
+LitGPT includes an experimental **Thunder** extension with enhanced FSDP:
+
+```bash
+litgpt pretrain tiny-llama-1.1b \
+  --devices 8 \
+  --num_nodes 1 \
+  --compiler thunder \
+  --strategy fsdp
+```
+
+### Thunder FSDP Configuration
+
+```python
+from extensions.thunder.pretrain import ThunderFSDPStrategy
+
+strategy = ThunderFSDPStrategy(
+    sharding_strategy="ZERO3",
+    bucketing_strategy="BLOCK",
+    state_dict_type="full",
+    jit=False,
+)
+```
+
+**Additional Parameters**:
+- `sharding_strategy`: `"ZERO3"` (full shard), `"ZERO2"` (grad/optim only)
+- `bucketing_strategy`: `"BLOCK"` (combine ops per block), `"LAYER"` (per layer), `"NONE"` (no bucketing)
+- `jit`: Whether to apply `thunder.jit(model)` for optimization
+- `executors`: Tuple of Thunder executors to enable
+
+**Bucketing Strategy**:
+- `"BLOCK"` (default): Combines collective operations for layer blocks → fewer communication calls
+- `"LAYER"`: Combines per layer class
+- `"NONE"`: No bucketing → more fine-grained but more overhead
+
+## Pretraining with FSDP
+
+### Single Node
+
+```bash
+litgpt pretrain tiny-llama-1.1b \
+  --devices 8 \
+  --train.global_batch_size 512 \
+  --train.micro_batch_size 8 \
+  --data Alpaca2k
+```
+
+**Memory calculation**:
+- TinyLlama 1.1B: ~4GB model + ~4GB gradients + ~8GB optimizer = 16GB per GPU without FSDP
+- With FSDP on 8 GPUs: 16GB / 8 = 2GB per GPU ✅ Fits easily
+
+### Multi-Node
+
+```bash
+# Launch on 4 nodes with 8 GPUs each (32 total)
+litgpt pretrain llama-2-7b \
+  --devices 8 \
+  --num_nodes 4 \
+  --train.global_batch_size 1024 \
+  --train.micro_batch_size 2 \
+  --data RedPajama
+```
+
+**Memory calculation**:
+- Llama 2 7B: ~28GB model + ~28GB gradients + ~56GB optimizer = 112GB total
+- With FSDP on 32 GPUs: 112GB / 32 = 3.5GB per GPU ✅
+
+## Fine-tuning with FSDP
+
+### LoRA Fine-tuning (Recommended)
+
+LoRA fine-tuning with FSDP for >7B models:
+
+```bash
+# Llama 2 70B LoRA on 8 GPUs
+litgpt finetune_lora meta-llama/Llama-2-70b-hf \
+  --devices 8 \
+  --data JSON \
+  --data.json_path data/alpaca.json \
+  --train.global_batch_size 16 \
+  --train.micro_batch_size 1 \
+  --lora_r 8
+```
+
+**Why LoRA with FSDP**:
+- Base model sharded with FSDP (memory efficient)
+- Only LoRA adapters trained (fast)
+- Best of both worlds for large models
+
+### Full Fine-tuning
+
+Full fine-tuning with FSDP:
+
+```bash
+# Llama 2 7B full fine-tune on 4 GPUs
+litgpt finetune_full meta-llama/Llama-2-7b-hf \
+  --devices 4 \
+  --data JSON \
+  --data.json_path data/alpaca.json \
+  --train.global_batch_size 16 \
+  --train.micro_batch_size 1 \
+  --train.learning_rate 3e-5
+```
+
+## Mixed Precision
+
+FSDP works with mixed precision for memory savings and speedup:
+
+```bash
+# BF16 mixed precision (recommended for A100/H100)
+litgpt pretrain tiny-llama-1.1b \
+  --devices 8 \
+  --precision bf16-mixed
+
+# FP16 mixed precision (V100 compatible)
+litgpt pretrain tiny-llama-1.1b \
+  --devices 8 \
+  --precision 16-mixed
+```
+
+**Precision options**:
+- `bf16-mixed`: BF16 for computation, FP32 for master weights (best for Ampere+)
+- `16-mixed`: FP16 for computation, FP32 for master weights (V100)
+- `32-true`: Full FP32 (debugging only, slow)
+
+## Gradient Accumulation
+
+Simulate larger batch sizes with gradient accumulation:
+
+```bash
+# Simulate global_batch_size=512 with micro_batch_size=2
+litgpt pretrain tiny-llama-1.1b \
+  --devices 8 \
+  --train.global_batch_size 512 \
+  --train.micro_batch_size 2
+# Accumulates over 512/(8*2) = 32 steps per optimizer update
+```
+
+**Formula**:
+```
+Gradient accumulation steps = global_batch_size / (devices × micro_batch_size)
+```
+
+## Memory Optimization
+
+### Out of Memory? Try These
+
+1. **Increase devices**:
+   ```bash
+   --devices 8  # Instead of 4
+   ```
+
+2. **Reduce micro batch size**:
+   ```bash
+   --train.micro_batch_size 1  # Instead of 2
+   ```
+
+3. **Lower precision**:
+   ```bash
+   --precision bf16-mixed  # Instead of 32-true
+   ```
+
+4. **Use FULL_SHARD**:
+   ```python
+   strategy = FSDPStrategy(
+       sharding_strategy="FULL_SHARD"  # Maximum memory savings
+   )
+   ```
+
+5. **Enable activation checkpointing** (implemented in model):
+   ```python
+   # Recomputes activations during backward pass
+   # Trades compute for memory
+   ```
+
+6. **Use QLoRA**:
+   ```bash
+   litgpt finetune_lora meta-llama/Llama-2-7b-hf \
+     --quantize bnb.nf4 \
+     --devices 1  # May not need FSDP with quantization
+   ```
+
+## Checkpointing
+
+### Save Checkpoints
+
+FSDP automatically handles checkpoint saving:
+
+```bash
+litgpt pretrain tiny-llama-1.1b \
+  --devices 8 \
+  --out_dir checkpoints/tinyllama-pretrain
+# Saves to: checkpoints/tinyllama-pretrain/final/lit_model.pth
+```
+
+With `state_dict_type="full"` (default), rank 0 assembles full model and saves single file.
+
+### Resume Training
+
+```bash
+litgpt pretrain tiny-llama-1.1b \
+  --devices 8 \
+  --resume checkpoints/tinyllama-pretrain/
+# Automatically loads latest checkpoint
+```
+
+### Convert to HuggingFace
+
+```bash
+python scripts/convert_lit_checkpoint.py \
+  --checkpoint_path checkpoints/tinyllama-pretrain/final/lit_model.pth \
+  --output_dir models/tinyllama-hf
+```
+
+## Performance Tuning
+
+### Communication Backends
+
+LitGPT uses NCCL for GPU communication:
+
+```bash
+# Default (NCCL auto-configured)
+litgpt pretrain tiny-llama-1.1b --devices 8
+
+# Explicit NCCL settings (advanced)
+NCCL_DEBUG=INFO \
+NCCL_IB_DISABLE=0 \
+litgpt pretrain tiny-llama-1.1b --devices 8
+```
+
+**NCCL Environment Variables**:
+- `NCCL_DEBUG=INFO`: Enable debug logging
+- `NCCL_IB_DISABLE=0`: Use InfiniBand (if available)
+- `NCCL_SOCKET_IFNAME=eth0`: Specify network interface
+
+### Multi-Node Setup
+
+**Option 1: SLURM**
+
+```bash
+#!/bin/bash
+#SBATCH --nodes=4
+#SBATCH --gpus-per-node=8
+#SBATCH --ntasks-per-node=1
+
+srun litgpt pretrain llama-2-7b \
+  --devices 8 \
+  --num_nodes 4 \
+  --data RedPajama
+```
+
+**Option 2: torchrun**
+
+```bash
+# On each node, run:
+torchrun \
+  --nproc_per_node=8 \
+  --nnodes=4 \
+  --node_rank=$NODE_RANK \
+  --master_addr=$MASTER_ADDR \
+  --master_port=29500 \
+  -m litgpt pretrain llama-2-7b
+```
+
+### Profiling
+
+Enable profiling to identify bottlenecks:
+
+```bash
+litgpt pretrain tiny-llama-1.1b \
+  --devices 8 \
+  --train.max_steps 100 \
+  --profile
+# Generates profiling report
+```
+
+## Example Configurations
+
+### Llama 2 7B on 4× A100 (40GB)
+
+```bash
+litgpt finetune_lora meta-llama/Llama-2-7b-hf \
+  --devices 4 \
+  --precision bf16-mixed \
+  --train.global_batch_size 64 \
+  --train.micro_batch_size 4 \
+  --train.max_seq_length 2048 \
+  --lora_r 8 \
+  --data JSON \
+  --data.json_path data/alpaca.json
+```
+
+**Memory per GPU**: ~20GB
+**Throughput**: ~5 samples/sec
+
+### Llama 2 70B on 8× A100 (80GB)
+
+```bash
+litgpt finetune_lora meta-llama/Llama-2-70b-hf \
+  --devices 8 \
+  --precision bf16-mixed \
+  --train.global_batch_size 32 \
+  --train.micro_batch_size 1 \
+  --train.max_seq_length 2048 \
+  --lora_r 8 \
+  --data JSON \
+  --data.json_path data/alpaca.json
+```
+
+**Memory per GPU**: ~70GB
+**Throughput**: ~1 sample/sec
+
+### Llama 3 405B on 64× H100 (80GB)
+
+```bash
+litgpt finetune_lora meta-llama/Llama-3.1-405B \
+  --devices 8 \
+  --num_nodes 8 \
+  --precision bf16-mixed \
+  --train.global_batch_size 128 \
+  --train.micro_batch_size 1 \
+  --train.max_seq_length 4096 \
+  --lora_r 16 \
+  --data JSON \
+  --data.json_path data/alpaca.json
+```
+
+**Memory per GPU**: ~60GB
+**Requires**: 64 H100 GPUs (8 nodes × 8 GPUs)
+
+## Troubleshooting
+
+### "CUDA out of memory"
+
+1. Reduce `micro_batch_size`
+2. Increase `devices` (more sharding)
+3. Lower `max_seq_length`
+4. Use `bf16-mixed` precision
+5. Try QLoRA (`--quantize bnb.nf4`)
+
+### "NCCL error" or Slow Communication
+
+1. Check network connectivity between nodes
+2. Enable InfiniBand: `NCCL_IB_DISABLE=0`
+3. Verify NCCL version: `python -c "import torch; print(torch.cuda.nccl.version())"`
+4. Test with NCCL tests: `$NCCL_HOME/build/all_reduce_perf -b 8 -e 128M`
+
+### Training Slower Than Expected
+
+1. Profile with `--profile`
+2. Check GPU utilization: `nvidia-smi dmon`
+3. Verify data loading isn't bottleneck
+4. Increase `micro_batch_size` if memory allows
+5. Use Thunder FSDP with bucketing
+
+## References
+
+- FSDP configuration: `litgpt/pretrain.py:setup()`
+- Thunder FSDP: `extensions/thunder/pretrain.py`
+- Memory optimization guide: `tutorials/oom.md`
+- Lightning Fabric docs: https://lightning.ai/docs/fabric/