diff --git a/optional-skills/mlops/accelerate/SKILL.md b/optional-skills/mlops/accelerate/SKILL.md index b77cf467735..e26f5503db4 100644 --- a/optional-skills/mlops/accelerate/SKILL.md +++ b/optional-skills/mlops/accelerate/SKILL.md @@ -1,7 +1,7 @@ --- name: huggingface-accelerate description: Run PyTorch training across GPUs with minimal changes. -version: 1.0.0 +version: 1.0.1 author: Orchestra Research license: MIT dependencies: [accelerate, torch, transformers] @@ -146,30 +146,35 @@ for batch in dataloader: ### Workflow 3: DeepSpeed ZeRO integration -**Enable DeepSpeed ZeRO-2**: +**Enable DeepSpeed ZeRO-2** (pass a `DeepSpeedPlugin`, not a raw dict): ```python -from accelerate import Accelerator +from accelerate import Accelerator, DeepSpeedPlugin + +deepspeed_plugin = DeepSpeedPlugin( + zero_stage=2, # ZeRO-2 + offload_optimizer_device="none", # or "cpu" to offload + gradient_accumulation_steps=4, +) accelerator = Accelerator( mixed_precision='bf16', - deepspeed_plugin={ - "zero_stage": 2, # ZeRO-2 - "offload_optimizer": False, - "gradient_accumulation_steps": 4 - } + deepspeed_plugin=deepspeed_plugin, # DeepSpeedPlugin instance (or dict[str, DeepSpeedPlugin]) ) # Same code as before! model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader) ``` -**Or via config**: -```bash -accelerate config -# Select: DeepSpeed → ZeRO-2 +**Or point at a full DeepSpeed JSON config via the plugin**: +```python +from accelerate import Accelerator, DeepSpeedPlugin + +# hf_ds_config accepts a path to a DeepSpeed config JSON (or a dict) +deepspeed_plugin = DeepSpeedPlugin(hf_ds_config="ds_config.json") +accelerator = Accelerator(mixed_precision='bf16', deepspeed_plugin=deepspeed_plugin) ``` -**deepspeed_config.json**: +**ds_config.json** (a raw DeepSpeed config — passed via the plugin, NOT via `--config_file`): ```json { "fp16": {"enabled": false}, @@ -183,9 +188,20 @@ accelerate config } ``` -**Launch**: +**Or via interactive config**: ```bash -accelerate launch --config_file deepspeed_config.json train.py +accelerate config +# Select: DeepSpeed → ZeRO-2 +# This writes an accelerate YAML config (default: ~/.cache/huggingface/accelerate/default_config.yaml) +``` + +**Launch** (`--config_file` expects an accelerate YAML, not a raw DeepSpeed JSON): +```bash +# Uses the default accelerate config written by `accelerate config` +accelerate launch train.py + +# Or point at a specific accelerate YAML +accelerate launch --config_file accelerate_deepspeed.yaml train.py ``` ### Workflow 4: FSDP (Fully Sharded Data Parallel) @@ -196,7 +212,7 @@ from accelerate import Accelerator, FullyShardedDataParallelPlugin fsdp_plugin = FullyShardedDataParallelPlugin( sharding_strategy="FULL_SHARD", # ZeRO-3 equivalent - auto_wrap_policy="TRANSFORMER_AUTO_WRAP", + auto_wrap_policy="transformer_based_wrap", # valid: transformer_based_wrap | size_based_wrap | no_wrap cpu_offload=False ) diff --git a/optional-skills/mlops/flash-attention/SKILL.md b/optional-skills/mlops/flash-attention/SKILL.md index 13eec2fa395..8f4374df688 100644 --- a/optional-skills/mlops/flash-attention/SKILL.md +++ b/optional-skills/mlops/flash-attention/SKILL.md @@ -1,7 +1,7 @@ --- name: optimizing-attention-flash description: Speed up long-sequence transformer training and inference. -version: 1.0.0 +version: 1.0.1 author: Orchestra Research license: MIT dependencies: [flash-attn, torch, transformers] @@ -82,13 +82,12 @@ import torch.nn.functional as F out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) ``` -Force Flash Attention backend: +Force Flash Attention backend (`torch.backends.cuda.sdp_kernel` is deprecated; use +`torch.nn.attention.sdpa_kernel` with `SDPBackend`): ```python -with torch.backends.cuda.sdp_kernel( - enable_flash=True, - enable_math=False, - enable_mem_efficient=False -): +from torch.nn.attention import SDPBackend, sdpa_kernel + +with sdpa_kernel(SDPBackend.FLASH_ATTENTION): out = F.scaled_dot_product_attention(q, k, v) ``` @@ -101,7 +100,8 @@ def test_attention(use_flash): q, k, v = [torch.randn(2, 8, 2048, 64, device='cuda', dtype=torch.float16) for _ in range(3)] if use_flash: - with torch.backends.cuda.sdp_kernel(enable_flash=True): + from torch.nn.attention import SDPBackend, sdpa_kernel + with sdpa_kernel(SDPBackend.FLASH_ATTENTION): return F.scaled_dot_product_attention(q, k, v) else: attn = (q @ k.transpose(-2, -1) / 8.0).softmax(dim=-1) @@ -230,14 +230,19 @@ print(f"Memory allocated: {torch.cuda.max_memory_allocated()/1e9:.2f}GB") ### Workflow 3: H100 FP8 optimization (FlashAttention-3) -For maximum performance on H100 GPUs. +For maximum performance on Hopper GPUs (H100). + +> **Important:** The pip package `flash-attn` (2.8.x) ships **FlashAttention-2 only** — it does +> **not** contain FA3 or FP8 H100 kernels, and `flash_attn_func` does **not** auto-use FP8. +> FlashAttention-3 is a separate **beta** build compiled from source from the repo's `hopper/` +> directory, exposed via the `flash_attn_interface` module. FA3 supports FP16/BF16 forward+backward +> and **FP8 forward only**. ``` FP8 Setup: -- [ ] Step 1: Verify H100 GPU available -- [ ] Step 2: Install flash-attn with FP8 support -- [ ] Step 3: Convert inputs to FP8 -- [ ] Step 4: Run with FP8 attention +- [ ] Step 1: Verify Hopper (H100) GPU available +- [ ] Step 2: Build & install FlashAttention-3 from source (hopper/) +- [ ] Step 3: Use the FA3 interface (FP8 forward) ``` **Step 1: Verify H100 GPU** @@ -247,36 +252,38 @@ nvidia-smi --query-gpu=name --format=csv # Should show "H100" or "H800" ``` -**Step 2: Install flash-attn with FP8 support** +**Step 2: Build & install FlashAttention-3 from source** + +FA3 is NOT included in `pip install flash-attn`. Build it from the `hopper/` subdirectory: ```bash -pip install flash-attn --no-build-isolation -# FP8 support included for H100 +git clone https://github.com/Dao-AILab/flash-attention.git +cd flash-attention/hopper +python setup.py install +# (compilation is heavy and requires a CUDA toolchain + Hopper GPU) ``` -**Step 3: Convert inputs to FP8** +**Step 3: Use the FA3 interface (FP8 forward)** + +FA3 exposes its own module `flash_attn_interface` (distinct from the FA2 `flash_attn`). +FP8 is a **forward-only** path and expects `float8_e4m3fn` inputs: ```python import torch +from flash_attn_interface import flash_attn_func # FA3 (hopper build), not `flash_attn` +# q, k, v: [batch, seqlen, nheads, headdim] q = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16) k = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16) v = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16) -# Convert to float8_e4m3 (FP8) +# FP8 forward (inference / forward-only): cast to float8_e4m3fn q_fp8 = q.to(torch.float8_e4m3fn) k_fp8 = k.to(torch.float8_e4m3fn) v_fp8 = v.to(torch.float8_e4m3fn) -``` -**Step 4: Run with FP8 attention** - -```python -from flash_attn import flash_attn_func - -# FlashAttention-3 automatically uses FP8 kernels on H100 -out = flash_attn_func(q_fp8, k_fp8, v_fp8) -# Result: ~1.2 PFLOPS, 1.5-2x faster than FP16 +out = flash_attn_func(q_fp8, k_fp8, v_fp8, causal=True) +# FP16/BF16 forward+backward is also supported by the FA3 interface. ``` ## When to use vs alternatives diff --git a/optional-skills/mlops/nemo-curator/SKILL.md b/optional-skills/mlops/nemo-curator/SKILL.md index 6045037491d..48cb211482f 100644 --- a/optional-skills/mlops/nemo-curator/SKILL.md +++ b/optional-skills/mlops/nemo-curator/SKILL.md @@ -1,7 +1,7 @@ --- name: nemo-curator description: "Curate LLM training data: dedupe, filter, PII redaction." -version: 1.0.0 +version: 1.0.1 author: Orchestra Research license: MIT dependencies: [nemo-curator, cudf, dask, rapids] @@ -40,41 +40,55 @@ NVIDIA's toolkit for preparing high-quality training data for LLMs. ### Installation ```bash +# NeMo Curator 1.x installs with uv. Extras use hyphens (PyPI-normalized): +# text-cuda12 / text-cpu (and image/video/audio/math variants), or `all`. + # Text curation (CUDA 12) -uv pip install "nemo-curator[text_cuda12]" +uv pip install "nemo-curator[text-cuda12]" # All modalities -uv pip install "nemo-curator[all_cuda12]" +uv pip install "nemo-curator[all]" -# CPU-only (slower) -uv pip install "nemo-curator[cpu]" +# CPU-only text (slower) +uv pip install "nemo-curator[text-cpu]" ``` ### Basic text curation pipeline +> **Major version rewrite (1.x):** NeMo Curator was rewritten around a **Ray-based +> pipeline/stage architecture**. The old `DocumentDataset` + `nemo_curator.modules.*` / +> `ScoreFilter` / `Modify` call-the-object-on-a-dataset API from 0.x is gone. In 1.x you +> compose `ProcessingStage`s into a `Pipeline` and run it with an executor. The exact +> stage/import surface differs per modality — treat the examples in this skill below as +> **conceptual** (0.x-style) and follow the current +> [quickstart](https://github.com/NVIDIA-NeMo/Curator/blob/main/tutorials/quickstart.py) +> and [text guide](https://docs.nvidia.com/nemo/curator/latest/get-started/text) for the +> exact 1.x APIs rather than copying imports verbatim. + +Shape of a 1.x pipeline (from the upstream quickstart): + ```python -from nemo_curator import ScoreFilter, Modify -from nemo_curator.datasets import DocumentDataset -import pandas as pd +from nemo_curator.pipeline import Pipeline +from nemo_curator.stages.base import ProcessingStage +from nemo_curator.stages.resources import Resources +from nemo_curator.backends.xenna import XennaExecutor +from nemo_curator.core.client import RayClient -# Load data -df = pd.DataFrame({"text": ["Good document", "Bad doc", "Excellent text"]}) -dataset = DocumentDataset(df) +# 1. Define/compose stages (load -> filter -> dedupe -> classify -> write). +# Each stage declares its own Resources (CPU cores, GPU memory, replicas). +pipeline = Pipeline(name="curation", stages=[...]) -# Quality filtering -def quality_score(doc): - return len(doc["text"].split()) > 5 # Filter short docs - -filtered = ScoreFilter(quality_score)(dataset) - -# Deduplication -from nemo_curator.modules import ExactDuplicates -deduped = ExactDuplicates()(filtered) - -# Save -deduped.to_parquet("curated_data/") +# 2. Run it with an executor (Ray-backed). +client = RayClient() +client.start() +pipeline.run(XennaExecutor()) +client.stop() ``` +The 0.x-style snippets in the sections that follow illustrate the *concepts* (quality +filtering, exact/fuzzy/semantic dedup, PII redaction, classifier filtering). For runnable +1.x code, map each concept onto the corresponding stage from the modality guide. + ## Data curation pipeline ### Stage 1: Quality filtering @@ -378,9 +392,9 @@ cluster.close() ## Resources -- **GitHub**: https://github.com/NVIDIA/NeMo-Curator ⭐ 500+ -- **Docs**: https://docs.nvidia.com/nemo-framework/user-guide/latest/datacuration/ -- **Version**: 0.4.0+ +- **GitHub**: https://github.com/NVIDIA-NeMo/Curator +- **Docs**: https://docs.nvidia.com/nemo/curator/latest/ +- **Version**: 1.2.0 (1.x is a Ray-based pipeline rewrite — see the quickstart before copying 0.x snippets) - **License**: Apache 2.0 diff --git a/optional-skills/mlops/saelens/SKILL.md b/optional-skills/mlops/saelens/SKILL.md index 9e0b894123f..813e5b3d58b 100644 --- a/optional-skills/mlops/saelens/SKILL.md +++ b/optional-skills/mlops/saelens/SKILL.md @@ -1,7 +1,7 @@ --- name: sparse-autoencoder-training description: Train sparse autoencoders to interpret model features. -version: 1.0.0 +version: 1.0.1 author: Orchestra Research license: MIT dependencies: [sae-lens>=6.0.0, transformer-lens>=2.0.0, torch>=2.0.0] @@ -78,11 +78,14 @@ from sae_lens import SAE # 1. Load model and pre-trained SAE model = HookedTransformer.from_pretrained("gpt2-small", device="cuda") -sae, cfg_dict, sparsity = SAE.from_pretrained( +# In sae-lens v6, SAE.from_pretrained() returns JUST the SAE (not a tuple). +sae = SAE.from_pretrained( release="gpt2-small-res-jb", sae_id="blocks.8.hook_resid_pre", device="cuda" ) +# If you also need the cfg dict and feature sparsity, use: +# sae, cfg_dict, sparsity = SAE.from_pretrained_with_cfg_and_sparsity(...) # 2. Get model activations tokens = model.to_tokens("The capital of France is Paris") @@ -124,24 +127,33 @@ reconstruction_error = (activations - reconstructed).norm() ### Step-by-Step ```python -from sae_lens import SAE, LanguageModelSAERunnerConfig, SAETrainingRunner +from sae_lens import ( + LanguageModelSAETrainingRunner, + LanguageModelSAERunnerConfig, + StandardTrainingSAEConfig, + LoggingConfig, +) -# 1. Configure training +# 1. Configure training (v6 uses a NESTED config: SAE-specific options live in a +# `sae=` sub-config, and logging options live in a `logger=` sub-config). +# Note: `architecture`, `d_sae`, `l1_coefficient` etc. are now on the SAE sub-config, +# and legacy flat options like `hook_layer`, `activation_fn`, `log_to_wandb` were removed. cfg = LanguageModelSAERunnerConfig( - # Model - model_name="gpt2-small", - hook_name="blocks.8.hook_resid_pre", - hook_layer=8, - d_in=768, # Model dimension + # SAE architecture + sparsity (nested) + sae=StandardTrainingSAEConfig( + d_in=768, # Model dimension + d_sae=768 * 8, # Expansion factor of 8 + l1_coefficient=8e-5, # Sparsity penalty + apply_b_dec_to_input=True, + normalize_activations="expected_average_only_in", + ), - # SAE architecture - architecture="standard", # or "gated", "topk" - d_sae=768 * 8, # Expansion factor of 8 - activation_fn="relu", + # Data-generating function (model + hook point) + model_name="gpt2-small", + hook_name="blocks.8.hook_resid_pre", # layer is inferred from hook_name (no hook_layer) # Training lr=4e-4, - l1_coefficient=8e-5, # Sparsity penalty l1_warm_up_steps=1000, train_batch_size_tokens=4096, training_tokens=100_000_000, @@ -150,9 +162,11 @@ cfg = LanguageModelSAERunnerConfig( dataset_path="monology/pile-uncopyrighted", context_size=128, - # Logging - log_to_wandb=True, - wandb_project="sae-training", + # Logging (nested) + logger=LoggingConfig( + log_to_wandb=True, + wandb_project="sae-training", + ), # Checkpointing checkpoint_path="checkpoints", @@ -160,7 +174,7 @@ cfg = LanguageModelSAERunnerConfig( ) # 2. Train -trainer = SAETrainingRunner(cfg) +trainer = LanguageModelSAETrainingRunner(cfg) # SAETrainingRunner still works as an alias sae = trainer.run() # 3. Evaluate @@ -168,6 +182,12 @@ print(f"L0 (avg active features): {trainer.metrics['l0']}") print(f"CE Loss Recovered: {trainer.metrics['ce_loss_score']}") ``` +> **v6 migration note:** For other SAE types swap the `sae=` sub-config — +> `GatedTrainingSAEConfig`, `TopKTrainingSAEConfig` (set `k` directly), or +> `JumpReLUTrainingSAEConfig` (uses `l0_coefficient`). Legacy flat options +> (`architecture`, `expansion_factor`, `hook_layer`, `activation_fn`/`activation_fn_kwargs`, +> `use_ghost_grads`, ghost grads, b_dec/decoder init options) were removed in v6. + ### Key Hyperparameters | Parameter | Typical Value | Effect | @@ -205,7 +225,7 @@ from sae_lens import SAE import torch model = HookedTransformer.from_pretrained("gpt2-small", device="cuda") -sae, _, _ = SAE.from_pretrained( +sae = SAE.from_pretrained( # v6 returns just the SAE release="gpt2-small-res-jb", sae_id="blocks.8.hook_resid_pre", device="cuda" @@ -279,47 +299,57 @@ for idx, val in zip(top_features.indices, top_features.values): ## Common Issues & Solutions +> All examples below use the v6 nested config: SAE-specific options go in the `sae=` +> sub-config (`StandardTrainingSAEConfig` / `TopKTrainingSAEConfig` / etc.), training +> knobs stay on the top-level `LanguageModelSAERunnerConfig`. + ### Issue: High dead feature ratio ```python -# WRONG: No warm-up, features die early +from sae_lens import LanguageModelSAERunnerConfig, StandardTrainingSAEConfig + +# WRONG: no warm-up, features die early cfg = LanguageModelSAERunnerConfig( - l1_coefficient=1e-4, + sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=1e-4), l1_warm_up_steps=0, # Bad! ) -# RIGHT: Warm-up L1 penalty +# RIGHT: warm up the L1 penalty (v6 removed ghost grads; warm-up is the lever now) cfg = LanguageModelSAERunnerConfig( - l1_coefficient=8e-5, + sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=8e-5), l1_warm_up_steps=1000, # Gradually increase - use_ghost_grads=True, # Revive dead features ) ``` ### Issue: Poor reconstruction (low CE recovery) ```python -# Reduce sparsity penalty +# Reduce sparsity penalty and/or add capacity (both on the SAE sub-config) cfg = LanguageModelSAERunnerConfig( - l1_coefficient=5e-5, # Lower = better reconstruction - d_sae=768 * 16, # More capacity + sae=StandardTrainingSAEConfig( + d_in=768, + d_sae=768 * 16, # More capacity + l1_coefficient=5e-5, # Lower = better reconstruction + ), ) ``` ### Issue: Features not interpretable ```python +from sae_lens import LanguageModelSAERunnerConfig, StandardTrainingSAEConfig, TopKTrainingSAEConfig + # Increase sparsity (higher L1) cfg = LanguageModelSAERunnerConfig( - l1_coefficient=1e-4, # Higher = sparser, more interpretable + sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=1e-4), ) -# Or use TopK architecture +# Or use a TopK SAE (k is set directly in v6, not via activation_fn_kwargs) cfg = LanguageModelSAERunnerConfig( - architecture="topk", - activation_fn_kwargs={"k": 50}, # Exactly 50 active features + sae=TopKTrainingSAEConfig(d_in=768, d_sae=768*8, k=50), # Exactly 50 active features ) ``` ### Issue: Memory errors during training ```python cfg = LanguageModelSAERunnerConfig( + sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=8e-5), train_batch_size_tokens=2048, # Reduce batch size store_batch_size_prompts=4, # Fewer prompts in buffer n_batches_in_buffer=8, # Smaller activation buffer @@ -341,8 +371,10 @@ Browse pre-trained SAE features at [neuronpedia.org](https://neuronpedia.org): | Class | Purpose | |-------|---------| | `SAE` | Sparse Autoencoder model | -| `LanguageModelSAERunnerConfig` | Training configuration | -| `SAETrainingRunner` | Training loop manager | +| `LanguageModelSAERunnerConfig` | Top-level training configuration (nests `sae=` and `logger=`) | +| `StandardTrainingSAEConfig` / `TopKTrainingSAEConfig` / `GatedTrainingSAEConfig` / `JumpReLUTrainingSAEConfig` | SAE-type-specific sub-configs (v6) | +| `LoggingConfig` | Logging/W&B sub-config (v6) | +| `LanguageModelSAETrainingRunner` | Training loop manager (alias: `SAETrainingRunner`) | | `ActivationsStore` | Activation collection and batching | | `HookedSAETransformer` | TransformerLens + SAE integration | @@ -381,10 +413,10 @@ For detailed API documentation, tutorials, and advanced usage, see the `referenc | **TopK** | Exactly K active features | Consistent sparsity | ```python -# TopK SAE (exactly 50 features active) +from sae_lens import LanguageModelSAERunnerConfig, TopKTrainingSAEConfig + +# TopK SAE (exactly 50 features active) — `k` is set on the SAE sub-config in v6 cfg = LanguageModelSAERunnerConfig( - architecture="topk", - activation_fn="topk", - activation_fn_kwargs={"k": 50}, + sae=TopKTrainingSAEConfig(d_in=768, d_sae=768*8, k=50), ) ``` diff --git a/optional-skills/mlops/tensorrt-llm/SKILL.md b/optional-skills/mlops/tensorrt-llm/SKILL.md index 6ef1f948d59..c5785c64dcf 100644 --- a/optional-skills/mlops/tensorrt-llm/SKILL.md +++ b/optional-skills/mlops/tensorrt-llm/SKILL.md @@ -1,7 +1,7 @@ --- name: tensorrt-llm description: High-throughput LLM inference on NVIDIA GPUs. -version: 1.0.0 +version: 1.0.1 author: Orchestra Research license: MIT dependencies: [tensorrt-llm, torch] @@ -40,13 +40,15 @@ NVIDIA's open-source library for optimizing LLM inference with state-of-the-art ### Installation ```bash -# Docker (recommended) -docker pull nvidia/tensorrt_llm:latest +# Docker (recommended) — images are on NGC (nvcr.io), not Docker Hub. +# Replace x.y.z with the desired version (e.g. 1.2.1). Browse tags on NGC: +# https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release/tags +docker pull nvcr.io/nvidia/tensorrt-llm/release:x.y.z -# pip install -pip install tensorrt_llm==1.2.0rc3 +# pip install (current stable GA) +pip install tensorrt_llm -# Requires CUDA 13.0.0, TensorRT 10.13.2, Python 3.10-3.12 +# Requires CUDA 13.2.1, TensorRT 10.x, Python 3.10-3.12 ``` ### Basic inference diff --git a/optional-skills/mlops/torchtitan/SKILL.md b/optional-skills/mlops/torchtitan/SKILL.md index 7da92663512..26fb68b9584 100644 --- a/optional-skills/mlops/torchtitan/SKILL.md +++ b/optional-skills/mlops/torchtitan/SKILL.md @@ -1,7 +1,7 @@ --- name: distributed-llm-pretraining-torchtitan description: Pretrain LLMs at scale with PyTorch 4D parallelism. -version: 1.0.0 +version: 1.0.1 author: Orchestra Research license: MIT dependencies: [torch>=2.6.0, torchtitan>=0.2.0, torchao>=0.5.0] @@ -37,7 +37,9 @@ python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets **Start training on 8 GPUs**: ```bash -CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh +# Configs are selected by name from the Python config registry +# (torchtitan/models/llama3/config_registry.py), not by TOML path +MODULE=llama3 CONFIG=llama3_8b ./run_train.sh ``` ## Common workflows @@ -65,10 +67,16 @@ python scripts/download_hf_assets.py \ **Step 2: Configure training** -Edit or create a TOML config file: +In torchtitan's current layout, run configs are defined in a Python **config registry** +(`torchtitan/models/llama3/config_registry.py`) and selected by name via `CONFIG=` +(or `--config `). To customize, register your own config in the registry, or override +individual fields on the command line (e.g. `--optimizer.lr 3e-4 --training.steps 1000`). + +The equivalent settings for an 8B run look like this (shown as fields; set them in the +registry entry or as `--section.key value` overrides): ```toml -# llama3_8b_custom.toml +# fields for a llama3 8B run (register in config_registry.py or pass as --overrides) [job] dump_folder = "./outputs" description = "Llama 3.1 8B training" @@ -108,13 +116,16 @@ interval = 500 **Step 3: Launch training** ```bash -# 8 GPUs on single node -CONFIG_FILE="./llama3_8b_custom.toml" ./run_train.sh +# 8 GPUs on single node (config selected by name from the registry) +MODULE=llama3 CONFIG=llama3_8b ./run_train.sh -# Or explicitly with torchrun +# Override individual fields on the command line +MODULE=llama3 CONFIG=llama3_8b ./run_train.sh --optimizer.lr 3e-4 --training.steps 1000 + +# Or explicitly with torchrun (run_train.sh wraps this) torchrun --nproc_per_node=8 \ -m torchtitan.train \ - --job.config_file ./llama3_8b_custom.toml + --module llama3 --config llama3_8b ``` **Step 4: Monitor and checkpoint** @@ -160,7 +171,7 @@ srun torchrun \ --rdzv_backend=c10d \ --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \ -m torchtitan.train \ - --job.config_file ./llama3_70b.toml + --module llama3 --config llama3_70b ``` **Step 3: Submit job** @@ -192,16 +203,28 @@ USE_CPP=0 pip install git+https://github.com/pytorch/ao.git **Step 2: Configure Float8** -Add to your TOML config: +In the current torchtitan, Float8 is applied at config time via the `quantization` +parameter in your `model_registry()` call inside the config registry (not via a +`[quantize.linear.float8]` TOML section). Add a `Float8LinearConverter.Config`: + +```python +# in torchtitan/models/llama3/config_registry.py (your model_registry(...) call) +from torchtitan.components.quantization import Float8LinearConverter + +model_spec = model_registry( + "8B", + quantization=[ + Float8LinearConverter.Config( + recipe_name="rowwise", # or "rowwise_with_gw_hp" + filter_fqns=["output"], # skip layers too small to benefit + model_compile_enabled=True, # requires torch.compile for competitive perf + ), + ], +) +``` + +Enable `torch.compile` in your run config too: ```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"] @@ -210,10 +233,8 @@ 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 +# Float8 config is baked into the registered config; just select it and enable compile +MODULE=llama3 CONFIG=llama3_8b ./run_train.sh --compile.enable ``` ### Workflow 4: 4D parallelism for 405B models @@ -229,7 +250,7 @@ CONFIG_FILE="./llama3_8b.toml" ./run_train.sh \ Required for consistent initialization across PP stages: ```bash -NGPU=1 CONFIG_FILE=./llama3_405b.toml ./run_train.sh \ +NGPU=1 MODULE=llama3 CONFIG=llama3_405b ./run_train.sh \ --checkpoint.enable \ --checkpoint.create_seed_checkpoint \ --parallelism.data_parallel_shard_degree 1 \ @@ -257,7 +278,7 @@ seq_len = 8192 # 64 nodes x 8 GPUs = 512 GPUs srun torchrun --nnodes=64 --nproc_per_node=8 \ -m torchtitan.train \ - --job.config_file ./llama3_405b.toml + --module llama3 --config llama3_405b ``` ## When to use vs alternatives @@ -304,10 +325,15 @@ 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"] +Float8 only benefits large GEMMs. Filter small layers via the converter's `filter_fqns`: +```python +from torchtitan.components.quantization import Float8LinearConverter + +Float8LinearConverter.Config( + # add "auto_filter_small_kn" to auto-skip layers too small to benefit + filter_fqns=["attention.wk", "attention.wv", "output", "auto_filter_small_kn"], + model_compile_enabled=True, +) ``` **Issue: Checkpoint loading fails after parallelism change** diff --git a/optional-skills/mlops/training/trl-fine-tuning/SKILL.md b/optional-skills/mlops/training/trl-fine-tuning/SKILL.md index 1fc6f6ccf58..8cc5ee99169 100644 --- a/optional-skills/mlops/training/trl-fine-tuning/SKILL.md +++ b/optional-skills/mlops/training/trl-fine-tuning/SKILL.md @@ -1,14 +1,14 @@ --- name: fine-tuning-with-trl -description: "TRL: SFT, DPO, PPO, GRPO, reward modeling for LLM RLHF." -version: 1.0.0 +description: "TRL: SFT, DPO, GRPO, RLOO reward modeling for LLM RLHF." +version: 1.0.1 author: Orchestra Research license: MIT dependencies: [trl, transformers, datasets, peft, accelerate, torch] platforms: [linux, macos, windows] metadata: hermes: - tags: [Post-Training, TRL, Reinforcement Learning, Fine-Tuning, SFT, DPO, PPO, GRPO, RLHF, Preference Alignment, HuggingFace] + tags: [Post-Training, TRL, Reinforcement Learning, Fine-Tuning, SFT, DPO, GRPO, RLOO, RLHF, Preference Alignment, HuggingFace] --- @@ -50,17 +50,23 @@ trainer.train() ## Common workflows -### Workflow 1: Full RLHF pipeline (SFT → Reward Model → PPO) +### Workflow 1: Full RLHF pipeline (SFT → Reward Model → RLOO) Complete pipeline from base model to human-aligned model. +> **Note (TRL 1.x):** PPO has been **removed** from TRL — `PPOTrainer`, `PPOConfig`, and +> `python -m trl.scripts.ppo` no longer exist. Use an online-RL trainer TRL still ships: +> **RLOO** (`RLOOTrainer` / `trl rloo`) is the closest drop-in for a reward-model-driven +> RLHF pipeline, and **GRPO** (`GRPOTrainer` / `trl grpo`, see Workflow 3) is the +> memory-efficient alternative. The step below uses RLOO. + Copy this checklist: ``` RLHF Training: - [ ] Step 1: Supervised fine-tuning (SFT) - [ ] Step 2: Train reward model -- [ ] Step 3: PPO reinforcement learning +- [ ] Step 3: RLOO reinforcement learning - [ ] Step 4: Evaluate aligned model ``` @@ -95,7 +101,7 @@ trainer = SFTTrainer( model=model, args=training_args, train_dataset=dataset, - tokenizer=tokenizer + processing_class=tokenizer ) trainer.train() trainer.save_model() @@ -138,19 +144,46 @@ trainer.train() trainer.save_model() ``` -**Step 3: PPO reinforcement learning** +**Step 3: RLOO reinforcement learning** -Optimize policy using reward model: +Optimize policy using the reward model. PPO was removed in TRL 1.x; use the RLOO CLI +(`trl rloo`) with the trained reward model passed via `--reward_model_name_or_path`: ```bash -python -m trl.scripts.ppo \ +trl rloo \ --model_name_or_path Qwen2.5-0.5B-SFT \ - --reward_model_path Qwen2.5-0.5B-Reward \ + --reward_model_name_or_path Qwen2.5-0.5B-Reward \ --dataset_name trl-internal-testing/descriptiveness-sentiment-trl-style \ - --output_dir Qwen2.5-0.5B-PPO \ + --output_dir Qwen2.5-0.5B-RLOO \ --learning_rate 3e-6 \ --per_device_train_batch_size 64 \ - --total_episodes 10000 + --num_generations 4 +``` + +Equivalent Python (`RLOOTrainer` / `RLOOConfig`): +```python +from trl import RLOOTrainer, RLOOConfig +from transformers import AutoModelForSequenceClassification, AutoTokenizer + +reward_model = AutoModelForSequenceClassification.from_pretrained( + "Qwen2.5-0.5B-Reward", num_labels=1 +) + +config = RLOOConfig( + output_dir="Qwen2.5-0.5B-RLOO", + per_device_train_batch_size=64, + learning_rate=3e-6, + num_generations=4, +) + +trainer = RLOOTrainer( + model="Qwen2.5-0.5B-SFT", + reward_funcs=reward_model, # a reward model (or a callable reward function) + args=config, + train_dataset=dataset, # prompt-only dataset + processing_class=tokenizer, +) +trainer.train() ``` **Step 4: Evaluate** @@ -159,7 +192,7 @@ python -m trl.scripts.ppo \ from transformers import pipeline # Load aligned model -generator = pipeline("text-generation", model="Qwen2.5-0.5B-PPO") +generator = pipeline("text-generation", model="Qwen2.5-0.5B-RLOO") # Test prompt = "Explain quantum computing to a 10-year-old" @@ -348,15 +381,15 @@ trl grpo \ **Use TRL when:** - Need to align model with human preferences - Have preference data (chosen/rejected pairs) -- Want to use reinforcement learning (PPO, GRPO) +- Want to use reinforcement learning (RLOO, GRPO) - Need reward model training - Doing RLHF (full pipeline) **Method selection**: - **SFT**: Have prompt-completion pairs, want basic instruction following - **DPO**: Have preferences, want simple alignment (no reward model needed) -- **PPO**: Have reward model, need maximum control over RL -- **GRPO**: Memory-constrained, want online RL +- **RLOO**: Have a reward model, want online RL (the reward-model-driven RLHF path; PPO was removed in TRL 1.x) +- **GRPO**: Memory-constrained, want online RL with reward functions - **Reward Model**: Building RLHF pipeline, need to score generations **Use alternatives instead:** @@ -411,13 +444,15 @@ print(dataset[0]) # Should have clear chosen > rejected ``` -**Issue: PPO training unstable** +**Issue: Online RL (RLOO/GRPO) training unstable** -Adjust KL coefficient: +Adjust the KL/beta regularization toward the reference policy: ```python -config = PPOConfig( - kl_coef=0.1, # Increase from 0.05 - cliprange=0.1 # Reduce from 0.2 +from trl import RLOOConfig + +config = RLOOConfig( + beta=0.05, # KL coefficient toward the reference model (increase for stability) + num_generations=4, # more samples per prompt = lower-variance advantage estimates ) ``` @@ -439,7 +474,7 @@ config = PPOConfig( - **VRAM**: Depends on model and method - SFT 7B: 16GB (with LoRA) - DPO 7B: 24GB (stores reference model) - - PPO 7B: 40GB (policy + reward model) + - RLOO 7B: 40GB (policy + reward model) - GRPO 7B: 24GB (more memory efficient) - **Multi-GPU**: Supported via `accelerate` - **Mixed precision**: BF16 recommended (A100/H100)