@r1z4x/tezgah 1.0.0

package/skills/ai-research/01-model-architecture/rwkv/references/state-management.md

@@ -0,0 +1,370 @@
+<!-- vendored from orchestra-research/AI-research-SKILLs@773a52944ba4747a18bd4ae9ade53fff041adcbc 01-model-architecture/rwkv/references/state-management.md; MIT (c) Orchestra Research -->
+# RWKV State Management
+
+## Understanding RWKV State
+
+Unlike Transformers with KV cache, RWKV maintains a **fixed-size recurrent state** that summarizes all previous context.
+
+### State Components
+
+```python
+state = {
+    'att_aa': torch.zeros(n_layers, d_model),  # Attention numerator accumulator
+    'att_ab': torch.zeros(n_layers, d_model),  # Attention denominator accumulator
+    'att_x_prev': torch.zeros(n_layers, d_model),  # Previous x for time-mixing
+    'ffn_x_prev': torch.zeros(n_layers, d_model)   # Previous x for channel-mixing
+}
+```
+
+**Total state size**: `4 × n_layers × d_model` parameters
+
+| Model | Layers | d_model | State Size |
+|-------|--------|---------|------------|
+| RWKV-169M | 12 | 768 | 37 KB |
+| RWKV-430M | 24 | 1024 | 98 KB |
+| RWKV-1.5B | 24 | 2048 | 196 KB |
+| RWKV-3B | 32 | 2560 | 327 KB |
+| RWKV-7B | 32 | 4096 | 524 KB |
+| RWKV-14B | 40 | 5120 | 819 KB |
+
+**Constant memory** regardless of context length!
+
+## State Initialization
+
+### Zero State (Default)
+
+```python
+from rwkv.model import RWKV
+
+model = RWKV(model='/path/to/RWKV-4-Pile-1B5', strategy='cuda fp16')
+
+# Start with zero state (no context)
+state = None
+out, state = model.forward(tokens, state)
+```
+
+### Warm State (Preloaded Context)
+
+```python
+# Load context once
+context = "The capital of France is Paris. The capital of Germany is Berlin."
+context_tokens = tokenizer.encode(context)
+
+# Process context to build state
+state = None
+for token in context_tokens:
+    _, state = model.forward([token], state)
+
+# Now use warm state for queries
+query = " The capital of Italy is"
+query_tokens = tokenizer.encode(query)
+out, state = model.forward(query_tokens, state)
+# Model "remembers" Paris and Berlin examples!
+```
+
+### Shared State (Multi-turn Conversations)
+
+```python
+# Conversation with persistent state
+state = None
+
+# Turn 1
+user1 = "My name is Alice."
+tokens1 = tokenizer.encode(user1)
+_, state = model.forward(tokens1, state)
+
+# Turn 2
+user2 = "What is my name?"
+tokens2 = tokenizer.encode(user2)
+response, state = model.forward(tokens2, state)
+# Response: "Alice" (state remembers!)
+```
+
+## State Update Rules
+
+### Time-Mixing State Update
+
+```python
+# Before processing token t
+att_aa_t = att_aa_{t-1}  # Previous numerator
+att_ab_t = att_ab_{t-1}  # Previous denominator
+
+# Compute WKV
+wkv_t = (exp(u) * k_t * v_t + att_aa_t) / (exp(u) * k_t + att_ab_t)
+
+# Update state for token t+1
+w = -exp(time_decay)  # Decay factor
+att_aa_{t+1} = exp(w) * att_aa_t + k_t * v_t
+att_ab_{t+1} = exp(w) * att_ab_t + k_t
+att_x_prev_{t+1} = x_t
+```
+
+**Effect of time_decay**:
+- **w = -0.01** (small decay): State decays slowly → long memory
+- **w = -5.0** (large decay): State decays quickly → short memory
+
+### Channel-Mixing State Update
+
+```python
+# Simply store previous x for next token
+ffn_x_prev_{t+1} = x_t
+```
+
+## State Serialization
+
+### Save/Load State (PyTorch)
+
+```python
+import torch
+
+# Save conversation state
+state_dict = {
+    'att_aa': state[0],
+    'att_ab': state[1],
+    'att_x_prev': state[2],
+    'ffn_x_prev': state[3]
+}
+torch.save(state_dict, 'conversation_123.pt')
+
+# Load state
+loaded = torch.load('conversation_123.pt')
+state = (loaded['att_aa'], loaded['att_ab'], loaded['att_x_prev'], loaded['ffn_x_prev'])
+
+# Continue conversation
+out, state = model.forward(new_tokens, state)
+```
+
+### State Compression (Optional)
+
+```python
+# FP16 state (half size)
+state_fp16 = tuple(s.half() for s in state)
+torch.save(state_fp16, 'state_compressed.pt')
+
+# Restore
+state = tuple(s.float() for s in torch.load('state_compressed.pt'))
+```
+
+## Multi-Session State Management
+
+### Session State Store
+
+```python
+class StateManager:
+    def __init__(self):
+        self.sessions = {}  # session_id -> state
+
+    def get_state(self, session_id):
+        return self.sessions.get(session_id, None)
+
+    def save_state(self, session_id, state):
+        self.sessions[session_id] = state
+
+    def clear_session(self, session_id):
+        if session_id in self.sessions:
+            del self.sessions[session_id]
+
+# Usage
+manager = StateManager()
+
+# User 1 conversation
+state1 = manager.get_state('user_1')
+out1, state1 = model.forward(tokens1, state1)
+manager.save_state('user_1', state1)
+
+# User 2 conversation (independent state)
+state2 = manager.get_state('user_2')
+out2, state2 = model.forward(tokens2, state2)
+manager.save_state('user_2', state2)
+```
+
+### State Expiration
+
+```python
+import time
+
+class StateManagerWithExpiry:
+    def __init__(self, expiry_seconds=3600):
+        self.sessions = {}  # session_id -> (state, timestamp)
+        self.expiry = expiry_seconds
+
+    def get_state(self, session_id):
+        if session_id in self.sessions:
+            state, timestamp = self.sessions[session_id]
+            if time.time() - timestamp < self.expiry:
+                return state
+            else:
+                del self.sessions[session_id]  # Expired
+        return None
+
+    def save_state(self, session_id, state):
+        self.sessions[session_id] = (state, time.time())
+```
+
+## State Interpolation
+
+### Blending States
+
+```python
+# Average two states (e.g., merging conversations)
+def blend_states(state1, state2, alpha=0.5):
+    """Blend state1 and state2 with weight alpha."""
+    return tuple(
+        alpha * s1 + (1 - alpha) * s2
+        for s1, s2 in zip(state1, state2)
+    )
+
+# Example: Blend Alice and Bob conversation contexts
+state_blended = blend_states(state_alice, state_bob, alpha=0.7)
+# 70% Alice context, 30% Bob context
+```
+
+### State Editing
+
+```python
+# Manually edit state (advanced)
+# Example: Reduce long-term memory influence
+
+def decay_state(state, decay_factor=0.5):
+    """Reduce state magnitude (forget older context)."""
+    att_aa, att_ab, att_x_prev, ffn_x_prev = state
+    return (
+        att_aa * decay_factor,
+        att_ab * decay_factor,
+        att_x_prev,  # Keep recent x
+        ffn_x_prev   # Keep recent x
+    )
+
+# Usage
+state = decay_state(state, decay_factor=0.3)  # Forget 70% of history
+```
+
+## Batch Inference with States
+
+### Independent Batch States
+
+```python
+# Each sequence in batch has separate state
+batch_size = 4
+states = [None] * batch_size
+
+for i, tokens in enumerate(batch_sequences):
+    out, states[i] = model.forward(tokens, states[i])
+```
+
+### Shared Prefix Optimization
+
+```python
+# All sequences share common prefix (e.g., system prompt)
+prefix = "You are a helpful assistant."
+prefix_tokens = tokenizer.encode(prefix)
+
+# Compute prefix state once
+prefix_state = None
+_, prefix_state = model.forward(prefix_tokens, None)
+
+# Clone prefix state for each sequence
+states = [prefix_state] * batch_size
+
+# Process user queries (independent)
+for i, user_query in enumerate(user_queries):
+    tokens = tokenizer.encode(user_query)
+    out, states[i] = model.forward(tokens, states[i])
+```
+
+## State Debugging
+
+### Inspect State Magnitudes
+
+```python
+def inspect_state(state):
+    """Print state statistics for debugging."""
+    att_aa, att_ab, att_x_prev, ffn_x_prev = state
+
+    print("State magnitudes:")
+    print(f"  att_aa: mean={att_aa.abs().mean():.4f}, max={att_aa.abs().max():.4f}")
+    print(f"  att_ab: mean={att_ab.abs().mean():.4f}, max={att_ab.abs().max():.4f}")
+    print(f"  att_x_prev: mean={att_x_prev.abs().mean():.4f}, max={att_x_prev.abs().max():.4f}")
+    print(f"  ffn_x_prev: mean={ffn_x_prev.abs().mean():.4f}, max={ffn_x_prev.abs().max():.4f}")
+
+# Usage
+inspect_state(state)
+```
+
+**Healthy ranges**:
+- `att_aa`, `att_ab`: 0.1 - 10.0 (if much larger, may overflow)
+- `att_x_prev`, `ffn_x_prev`: Similar to input embedding magnitude
+
+### State Divergence Check
+
+```python
+def state_distance(state1, state2):
+    """Compute L2 distance between two states."""
+    return sum(
+        torch.dist(s1, s2).item()
+        for s1, s2 in zip(state1, state2)
+    )
+
+# Example: Check if states diverged
+distance = state_distance(state_alice, state_bob)
+print(f"State distance: {distance:.2f}")
+# Large distance → very different contexts
+```
+
+## Numerical Stability Considerations
+
+### Overflow Prevention
+
+```python
+# Issue: att_aa, att_ab can grow unbounded
+# If att_aa > 1e10, numerical precision issues
+
+# Solution 1: Periodic normalization
+if att_aa.abs().max() > 1e6:
+    scale = att_aa.abs().max()
+    att_aa = att_aa / scale
+    att_ab = att_ab / scale
+```
+
+### Underflow Prevention
+
+```python
+# Issue: With large negative time_decay, state can underflow to 0
+
+# Solution: Clip time_decay
+time_decay = torch.clamp(time_decay, min=-8.0, max=-0.1)
+# Ensures state doesn't decay too fast
+```
+
+## State vs KV Cache Comparison
+
+### Memory Usage (8K context)
+
+| Model Type | Model Size | KV Cache Size | RWKV State Size |
+|------------|------------|---------------|-----------------|
+| Transformer | 1.3B | 4.1 GB | - |
+| **RWKV** | **1.5B** | **-** | **196 KB** |
+| Transformer | 7B | 21.3 GB | - |
+| **RWKV** | **7B** | **-** | **524 KB** |
+
+**RWKV advantage**: 10,000× smaller than KV cache!
+
+### Information Retention
+
+**KV Cache (Transformer)**:
+- Perfect: Stores all previous keys and values
+- Retrieval: Exact attention to any previous token
+- Cost: O(n) memory growth
+
+**RWKV State**:
+- Lossy: Compressed representation of history
+- Retrieval: Weighted blend of previous tokens (decay-based)
+- Cost: O(1) constant memory
+
+**Trade-off**: RWKV sacrifices perfect recall for constant memory
+
+## Resources
+
+- State management examples: https://github.com/BlinkDL/ChatRWKV
+- Wiki: https://wiki.rwkv.com/state-management
+- Discord: https://discord.gg/bDSBUMeFpc (RWKV community)

package/skills/ai-research/01-model-architecture/torchtitan/SKILL.md

@@ -0,0 +1,358 @@
+---
+name: distributed-llm-pretraining-torchtitan
+description: Provides PyTorch-native distributed LLM pretraining using torchtitan with 4D parallelism (FSDP2, TP, PP, CP). Use when pretraining Llama 3.1, DeepSeek V3, or custom models at scale from 8 to 512+ GPUs with Float8, torch.compile, and distributed checkpointing.
+version: 1.0.0
+author: Orchestra Research
+license: MIT
+tags: [Model Architecture, Distributed Training, TorchTitan, FSDP2, Tensor Parallel, Pipeline Parallel, Context Parallel, Float8, Llama, Pretraining]
+dependencies: [torch>=2.6.0, torchtitan>=0.2.0, torchao>=0.5.0]
+---
+
+# TorchTitan - PyTorch Native Distributed LLM Pretraining
+
+## Quick start
+
+TorchTitan is PyTorch's official platform for large-scale LLM pretraining with composable 4D parallelism (FSDP2, TP, PP, CP), achieving 65%+ speedups over baselines on H100 GPUs.
+
+**Installation**:
+```bash
+# From PyPI (stable)
+pip install torchtitan
+
+# From source (latest features, requires PyTorch nightly)
+git clone https://github.com/pytorch/torchtitan
+cd torchtitan
+pip install -r requirements.txt
+```
+
+**Download tokenizer**:
+```bash
+# Get HF token from https://huggingface.co/settings/tokens
+python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...
+```
+
+**Start training on 8 GPUs**:
+```bash
+CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh
+```
+
+## Common workflows
+
+### Workflow 1: Pretrain Llama 3.1 8B on single node
+
+Copy this checklist:
+
+```
+Single Node Pretraining:
+- [ ] Step 1: Download tokenizer
+- [ ] Step 2: Configure training
+- [ ] Step 3: Launch training
+- [ ] Step 4: Monitor and checkpoint
+```
+
+**Step 1: Download tokenizer**
+
+```bash
+python scripts/download_hf_assets.py \
+  --repo_id meta-llama/Llama-3.1-8B \
+  --assets tokenizer \
+  --hf_token=YOUR_HF_TOKEN
+```
+
+**Step 2: Configure training**
+
+Edit or create a TOML config file:
+
+```toml
+# llama3_8b_custom.toml
+[job]
+dump_folder = "./outputs"
+description = "Llama 3.1 8B training"
+
+[model]
+name = "llama3"
+flavor = "8B"
+hf_assets_path = "./assets/hf/Llama-3.1-8B"
+
+[optimizer]
+name = "AdamW"
+lr = 3e-4
+
+[lr_scheduler]
+warmup_steps = 200
+
+[training]
+local_batch_size = 2
+seq_len = 8192
+max_norm = 1.0
+steps = 1000
+dataset = "c4"
+
+[parallelism]
+data_parallel_shard_degree = -1  # Use all GPUs for FSDP
+
+[activation_checkpoint]
+mode = "selective"
+selective_ac_option = "op"
+
+[checkpoint]
+enable = true
+folder = "checkpoint"
+interval = 500
+```
+
+**Step 3: Launch training**
+
+```bash
+# 8 GPUs on single node
+CONFIG_FILE="./llama3_8b_custom.toml" ./run_train.sh
+
+# Or explicitly with torchrun
+torchrun --nproc_per_node=8 \
+  -m torchtitan.train \
+  --job.config_file ./llama3_8b_custom.toml
+```
+
+**Step 4: Monitor and checkpoint**
+
+TensorBoard logs are saved to `./outputs/tb/`:
+```bash
+tensorboard --logdir ./outputs/tb
+```
+
+### Workflow 2: Multi-node training with SLURM
+
+```
+Multi-Node Training:
+- [ ] Step 1: Configure parallelism for scale
+- [ ] Step 2: Set up SLURM script
+- [ ] Step 3: Submit job
+- [ ] Step 4: Resume from checkpoint
+```
+
+**Step 1: Configure parallelism for scale**
+
+For 70B model on 256 GPUs (32 nodes):
+```toml
+[parallelism]
+data_parallel_shard_degree = 32  # FSDP across 32 ranks
+tensor_parallel_degree = 8        # TP within node
+pipeline_parallel_degree = 1      # No PP for 70B
+context_parallel_degree = 1       # Increase for long sequences
+```
+
+**Step 2: Set up SLURM script**
+
+```bash
+#!/bin/bash
+#SBATCH --job-name=llama70b
+#SBATCH --nodes=32
+#SBATCH --ntasks-per-node=8
+#SBATCH --gpus-per-node=8
+
+srun torchrun \
+  --nnodes=32 \
+  --nproc_per_node=8 \
+  --rdzv_backend=c10d \
+  --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
+  -m torchtitan.train \
+  --job.config_file ./llama3_70b.toml
+```
+
+**Step 3: Submit job**
+
+```bash
+sbatch multinode_trainer.slurm
+```
+
+**Step 4: Resume from checkpoint**
+
+Training auto-resumes if checkpoint exists in configured folder.
+
+### Workflow 3: Enable Float8 training for H100s
+
+Float8 provides 30-50% speedup on H100 GPUs.
+
+```
+Float8 Training:
+- [ ] Step 1: Install torchao
+- [ ] Step 2: Configure Float8
+- [ ] Step 3: Launch with compile
+```
+
+**Step 1: Install torchao**
+
+```bash
+USE_CPP=0 pip install git+https://github.com/pytorch/ao.git
+```
+
+**Step 2: Configure Float8**
+
+Add to your TOML config:
+```toml
+[model]
+converters = ["quantize.linear.float8"]
+
+[quantize.linear.float8]
+enable_fsdp_float8_all_gather = true
+precompute_float8_dynamic_scale_for_fsdp = true
+filter_fqns = ["output"]  # Exclude output layer
+
+[compile]
+enable = true
+components = ["model", "loss"]
+```
+
+**Step 3: Launch with compile**
+
+```bash
+CONFIG_FILE="./llama3_8b.toml" ./run_train.sh \
+  --model.converters="quantize.linear.float8" \
+  --quantize.linear.float8.enable_fsdp_float8_all_gather \
+  --compile.enable
+```
+
+### Workflow 4: 4D parallelism for 405B models
+
+```
+4D Parallelism (FSDP + TP + PP + CP):
+- [ ] Step 1: Create seed checkpoint
+- [ ] Step 2: Configure 4D parallelism
+- [ ] Step 3: Launch on 512 GPUs
+```
+
+**Step 1: Create seed checkpoint**
+
+Required for consistent initialization across PP stages:
+```bash
+NGPU=1 CONFIG_FILE=./llama3_405b.toml ./run_train.sh \
+  --checkpoint.enable \
+  --checkpoint.create_seed_checkpoint \
+  --parallelism.data_parallel_shard_degree 1 \
+  --parallelism.tensor_parallel_degree 1 \
+  --parallelism.pipeline_parallel_degree 1
+```
+
+**Step 2: Configure 4D parallelism**
+
+```toml
+[parallelism]
+data_parallel_shard_degree = 8   # FSDP
+tensor_parallel_degree = 8       # TP within node
+pipeline_parallel_degree = 8     # PP across nodes
+context_parallel_degree = 1      # CP for long sequences
+
+[training]
+local_batch_size = 32
+seq_len = 8192
+```
+
+**Step 3: Launch on 512 GPUs**
+
+```bash
+# 64 nodes x 8 GPUs = 512 GPUs
+srun torchrun --nnodes=64 --nproc_per_node=8 \
+  -m torchtitan.train \
+  --job.config_file ./llama3_405b.toml
+```
+
+## When to use vs alternatives
+
+**Use TorchTitan when:**
+- Pretraining LLMs from scratch (8B to 405B+)
+- Need PyTorch-native solution without third-party dependencies
+- Require composable 4D parallelism (FSDP2, TP, PP, CP)
+- Training on H100s with Float8 support
+- Want interoperable checkpoints with torchtune/HuggingFace
+
+**Use alternatives instead:**
+- **Megatron-LM**: Maximum performance for NVIDIA-only deployments
+- **DeepSpeed**: Broader ZeRO optimization ecosystem, inference support
+- **Axolotl/TRL**: Fine-tuning rather than pretraining
+- **LitGPT**: Educational, smaller-scale training
+
+## Common issues
+
+**Issue: Out of memory on large models**
+
+Enable activation checkpointing and reduce batch size:
+```toml
+[activation_checkpoint]
+mode = "full"  # Instead of "selective"
+
+[training]
+local_batch_size = 1
+```
+
+Or use gradient accumulation:
+```toml
+[training]
+local_batch_size = 1
+global_batch_size = 32  # Accumulates gradients
+```
+
+**Issue: TP causes high memory with async collectives**
+
+Set environment variable:
+```bash
+export TORCH_NCCL_AVOID_RECORD_STREAMS=1
+```
+
+**Issue: Float8 training not faster**
+
+Float8 only benefits large GEMMs. Filter small layers:
+```toml
+[quantize.linear.float8]
+filter_fqns = ["attention.wk", "attention.wv", "output", "auto_filter_small_kn"]
+```
+
+**Issue: Checkpoint loading fails after parallelism change**
+
+Use DCP's resharding capability:
+```bash
+# Convert sharded checkpoint to single file
+python -m torch.distributed.checkpoint.format_utils \
+  dcp_to_torch checkpoint/step-1000 checkpoint.pt
+```
+
+**Issue: Pipeline parallelism initialization**
+
+Create seed checkpoint first (see Workflow 4, Step 1).
+
+## Supported models
+
+| Model | Sizes | Status |
+|-------|-------|--------|
+| Llama 3.1 | 8B, 70B, 405B | Production |
+| Llama 4 | Various | Experimental |
+| DeepSeek V3 | 16B, 236B, 671B (MoE) | Experimental |
+| GPT-OSS | 20B, 120B (MoE) | Experimental |
+| Qwen 3 | Various | Experimental |
+| Flux | Diffusion | Experimental |
+
+## Performance benchmarks (H100)
+
+| Model | GPUs | Parallelism | TPS/GPU | Techniques |
+|-------|------|-------------|---------|------------|
+| Llama 8B | 8 | FSDP | 5,762 | Baseline |
+| Llama 8B | 8 | FSDP+compile+FP8 | 8,532 | +48% |
+| Llama 70B | 256 | FSDP+TP+AsyncTP | 876 | 2D parallel |
+| Llama 405B | 512 | FSDP+TP+PP | 128 | 3D parallel |
+
+## Advanced topics
+
+**FSDP2 configuration**: See [references/fsdp.md](references/fsdp.md) for detailed FSDP2 vs FSDP1 comparison and ZeRO equivalents.
+
+**Float8 training**: See [references/float8.md](references/float8.md) for tensorwise vs rowwise scaling recipes.
+
+**Checkpointing**: See [references/checkpoint.md](references/checkpoint.md) for HuggingFace conversion and async checkpointing.
+
+**Adding custom models**: See [references/custom-models.md](references/custom-models.md) for TrainSpec protocol.
+
+## Resources
+
+- GitHub: https://github.com/pytorch/torchtitan
+- Paper: https://arxiv.org/abs/2410.06511
+- ICLR 2025: https://iclr.cc/virtual/2025/poster/29620
+- PyTorch Forum: https://discuss.pytorch.org/c/distributed/torchtitan/44
+