diff --git a/cookbooks/cosmos3/generator/audiovisual/finetune/README.md b/cookbooks/cosmos3/generator/audiovisual/finetune/README.md index 3d2d7ee7..3962de4a 100644 --- a/cookbooks/cosmos3/generator/audiovisual/finetune/README.md +++ b/cookbooks/cosmos3/generator/audiovisual/finetune/README.md @@ -34,6 +34,38 @@ Paths are fixed at the top of each script (under this git-ignored folder) — ed These recipes default to 8 GPUs. On a 4-GPU node (e.g. GB200×4), set `--nproc_per_node=4` on the `torchrun` line in the launch script. +## LoRA fine-tuning + +LoRA freezes the backbone and trains small low-rank adapters on the generation-pathway attention projections, so optimizer state is adapter-sized rather than backbone-sized — that is what lets the 32B Super tier fit a node budget a full fine-tune of that size would not. + +Enabling it is a **TOML-level switch**: no code change and no separate LoRA experiment to register. `[job].experiment` stays whatever registered experiment you started from (`vision_sft_nano`, `vision_sft_edge`, …); change `[job].name` so the run gets its own output directory. + +**[`launch_sft_vision_super.sh`](launch_sft_vision_super.sh) is already a LoRA recipe** — read it and its TOML, [`toml/sft_config/vision_sft_super.toml`](toml/sft_config/vision_sft_super.toml), alongside this section for a known-good set of values. To convert another recipe, copy its TOML and touch these four places: + +| TOML section | Change | Why | +| --- | --- | --- | +| `[model]` | `lora_enabled = true`, plus `lora_rank` / `lora_alpha` / `lora_target_modules` | injects the adapters before FSDP wraps the network | +| `[optimizer]` | `keys_to_select = ["lora_"]`, and raise `lr` | trains the adapters only; they start at zero and are a tiny fraction of the parameters, so they tolerate — and want — a larger LR than a full fine-tune | +| `[checkpoint]` | add `"lora_"` to `keys_to_skip_loading` | the base checkpoint has no adapter tensors; let them initialize fresh instead of failing the load | +| `[model.ema]` / `[model.compile]` | disable both | EMA over a frozen backbone is wasted memory, and LoRA injection subclasses `nn.Linear`, so leave `torch.compile` off | + +The `[model]` block ends up looking like the following — the numbers are the Super recipe's, shown as an example to start from rather than a required setting: + +```toml +lora_enabled = true +lora_rank = 16 # e.g. — larger rank buys capacity at the cost of more trainable params +lora_alpha = 32 # e.g. — commonly around 2x rank +lora_target_modules = "q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen" +``` + +`lora_target_modules` matches a bare name against exact child-module names, and any target containing a `.` against the full module path. The four MoE-gen projections above are the generation pathway and exist in every MoT block on all tiers (dense-FFN Edge included), so that line carries over unchanged — and because `_moe_gen` is part of the leaf name itself, they can't collide with the understanding tower. + +Extending beyond attention needs more care: the gen-path FFN is `mlp_moe_gen`, but its children are named exactly like the understanding tower's (`mlp.up_proj`, …), so a bare `up_proj` would adapt both. Write those targets path-qualified instead — `mlp_moe_gen.up_proj`, and so on — and check the module tree first, since the FFN's children differ between the dense and MoE tiers. + +For the LR, take the recipe's full-FT value as the starting point and go up — for example, the Super LoRA recipe runs `5e-4` where the full fine-tunes use `1e-4`. + +Finally, copy the matching launch shell and point `--sft-toml` at the new TOML. Everything else in the recipe (data, VAE, checkpoint conversion), the outputs, and the [Hugging Face export](#export-to-hugging-face-safetensors) are unchanged. + ## Outputs Training writes to `outputs/train////`: diff --git a/cookbooks/cosmos3/reasoner/finetune/README.md b/cookbooks/cosmos3/reasoner/finetune/README.md index 8fa23621..132a4a9a 100644 --- a/cookbooks/cosmos3/reasoner/finetune/README.md +++ b/cookbooks/cosmos3/reasoner/finetune/README.md @@ -37,6 +37,38 @@ The VideoPhy-2 download/convert steps are skipped once their outputs exist (Edge These recipes default to 8 GPUs. On a 4-GPU node (e.g. GB200×4), set `--nproc_per_node=4` on the `torchrun` line in the launch script. +## LoRA fine-tuning + +LoRA freezes the reasoner backbone and trains small low-rank adapters on the LLM attention projections, so only the adapters carry optimizer state — which is what makes the 32B Super tier comfortable on a 4-GPU allocation. + +Every recipe above is a full fine-tune, but LoRA is a **TOML-level switch**: no code change and no separate LoRA experiment to register. Keep `[job].experiment` on the recipe you started from (`videophy2_sft_nano` / `_super` / `_edge`, or `pre_exp012_llava_ov` for the LLaVA-OneVision recipe) and change `[job].name` so the run gets its own output directory. The generator cookbook's [`launch_sft_vision_super.sh`](../../generator/audiovisual/finetune/launch_sft_vision_super.sh) is a ready-made LoRA recipe driven by the same mechanism on the `vfm` side — worth reading alongside. + +Copy the recipe's TOML and touch two places: + +| TOML section | Change | Why | +| --- | --- | --- | +| `[model]` | `lora_enabled = true`, plus `lora_rank` / `lora_alpha` / `lora_target_modules` | injects the adapters before FSDP wraps the network | +| `[optimizer]` | `keys_to_select = ["lora_"]`, and raise `lr` | trains the adapters only; they start at zero and are a tiny fraction of the parameters, so they take a larger LR than the full-FT recipe's | + +For the reasoner backbones the targets are the LLM attention projections — the rank/alpha below are an example to start from, not a required setting: + +```toml +lora_enabled = true +lora_rank = 16 # e.g. — larger rank buys capacity at the cost of more trainable params +lora_alpha = 32 # e.g. — commonly around 2x rank +lora_target_modules = "q_proj,k_proj,v_proj,o_proj" +``` + +Names are matched against the *live* module tree, not checkpoint keys, and these four cover both Qwen3-VL (Nano/Super) and the Cosmos3-Edge reasoner; adding the MLP projections (`gate_proj,up_proj,down_proj` on Qwen3-VL, `up_proj,down_proj` on Edge) is the usual way to buy more capacity. For the LR, start from the recipe's full-FT value and go up — for example, ~5× (`1e-6` → `5e-6`). + +**On Cosmos3-Edge add one more line:** `lora_exclude_path_regex = "^model\\.visual\\."`. Its SigLIP2 vision tower names three of its four projections `q_proj`/`k_proj`/`v_proj` too, so name matching alone would also adapt the tower the recipe deliberately freezes — and the run would look healthy while training the wrong subnetwork. + +Nothing else *has* to change: `[model.ema].enabled` and `[model.compile].enabled` are already `false` in these recipes (LoRA injection subclasses `nn.Linear`, so keep compile off), and `[checkpoint].keys_to_skip_loading` stays `[]` — the VLM backbone loads from HF safetensors, not from a DCP checkpoint that could carry stale adapter keys. Separately from LoRA, the VideoPhy-2 recipes are short smoke runs (`[trainer].max_iter = 50`), so a real run also wants a higher `max_iter` with `[scheduler].cycle_lengths` matched to it. + +Then copy the matching launch shell and point `--sft-toml` at the new TOML; data prep and checkpoint preparation are unchanged. + +The VideoPhy-2 recipes enable the HF export callback, so every checkpoint save also writes an HF snapshot next to the DCP checkpoint. That snapshot is an ordinary HF checkpoint even for a LoRA run — the callback merges each adapter into its base weight (`W + (alpha/r) · B·A`), so it carries no `lora_*` keys and loads exactly like a full fine-tune's. Gathering the whole backbone on rank 0 per save is wasted work on a short convergence run, so pass `-- checkpoint.hf_export.enabled=false` on the `torchrun` line while you're only watching loss curves. + ## Outputs Training writes to `outputs/train////`: