diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index dd01a287..be7450e3 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -337,6 +337,16 @@ jobs: uv run --all-extras --group=cu128-train torchrun --nproc_per_node=4 -m pytest -v \ cosmos_framework/model/generator/mot/context_parallel_test.py -o addopts= + # HFExportCallback folds LoRA adapters into the base weights and gathers + # lora_A / lora_B from inside the base weight's loop iteration, so the + # ordering is only exercised once the adapters are DTensors. World size is + # fixed at 2; the test skips itself under any other. + - name: Distributed unit tests - hf_export LoRA merge (torchrun, 2 ranks) + run: | + export LD_LIBRARY_PATH= + uv run --all-extras --group=cu128-train torchrun --nproc_per_node=2 -m pytest -v \ + cosmos_framework/callbacks/hf_export_fsdp_test.py -o addopts= + # Clear everything the suite writes into the working tree (all gitignored # scratch): pytest tmp dirs (DCP checkpoint, logs), the script-test # `outputs/` dir, any `examples/checkpoints`, and the `schemas/` dir from diff --git a/cosmos_framework/callbacks/hf_export.py b/cosmos_framework/callbacks/hf_export.py index 88c9069e..86ca9c55 100644 --- a/cosmos_framework/callbacks/hf_export.py +++ b/cosmos_framework/callbacks/hf_export.py @@ -12,6 +12,11 @@ - Worker exceptions are stored in ``_worker_exception`` and re-raised on the next checkpoint or at train end, so failures are never silently swallowed. - Controlled entirely via ``config.checkpoint.hf_export`` (HFExportConfig). +- LoRA runs export a MERGED checkpoint: each ``LoraInjectedLinear`` contributes a + single ``.weight`` equal to ``W + (alpha / r) * B @ A``, and no ``lora_*`` + keys are written. The export is therefore a plain HF checkpoint that + ``from_pretrained`` (and ``eval_videophy2``) loads with the adapter's effect + already in the weights, exactly like a full fine-tune's export. Phase 2+ note ------------- @@ -205,6 +210,60 @@ def on_train_end(self, model: Any, iteration: int = 0) -> None: # Internal helpers # ------------------------------------------------------------------ + # Path segments torch.compile and gradient checkpointing insert into the + # module tree. They are not part of the HF-native name. + _WRAPPER_SEGMENTS: frozenset[str] = frozenset({"_orig_mod", "_checkpoint_wrapped_module"}) + + @classmethod + def _strip_wrapper_prefixes(cls, name: str) -> str: + """Drop wrapper segments from a parameter or module path. + + Dropping whole dot-separated segments rather than substrings matters for + module paths: a wrapped module's own path *ends* with the wrapper segment + (``layer._checkpoint_wrapped_module``) and has no trailing dot to match + on, so substring removal would leave it alone and + :meth:`_lora_merge_plan` would key its adapters off a name that no + stripped parameter ever produces — a silently unmerged export. + """ + return ".".join(seg for seg in name.split(".") if seg not in cls._WRAPPER_SEGMENTS) + + @staticmethod + def _lora_merge_plan(root: torch.nn.Module) -> tuple[dict[str, Any], set[str]]: + """Locate every LoRA-adapted linear and the adapter keys it owns. + + Returns ``(merge_targets, adapter_keys)`` where ``merge_targets`` maps a + base-weight parameter name to the ``LoraInjectedLinear`` holding it, and + ``adapter_keys`` is the set of ``lora_A`` / ``lora_B`` parameter names + that must NOT be written to the export. Both use post-strip names so + they match what :meth:`_gather_weights` computes. + + Empty on a full fine-tune, which is what keeps that path untouched. + """ + # Deferred: cosmos_framework.utils.generator.lora is only needed when a + # LoRA run reaches export, and hf_export is imported from config land. + from cosmos_framework.utils.generator.lora import LoraInjectedLinear + + merge_targets: dict[str, Any] = {} + adapter_keys: set[str] = set() + for module_name, module in root.named_modules(): + if not isinstance(module, LoraInjectedLinear): + continue + path = HFExportCallback._strip_wrapper_prefixes(module_name) + # An adapted linear at the tree root (or under nothing but wrappers) + # strips to "", and its parameters are plain "weight" / "lora_A.weight". + prefix = f"{path}." if path else "" + merge_targets[f"{prefix}weight"] = module + adapter_keys.add(f"{prefix}lora_A.weight") + adapter_keys.add(f"{prefix}lora_B.weight") + return merge_targets, adapter_keys + + @staticmethod + def _gather_full(param: torch.Tensor) -> torch.Tensor: + """All-gather a sharded parameter. Collective — every rank must call it.""" + if isinstance(param, torch.distributed.tensor.DTensor): + param = param.full_tensor() + return param.detach() + def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], dict[str, str], int]: """Iterate model parameters, all-gather DTensor shards, and build CPU chunks. @@ -212,17 +271,28 @@ def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], di ``cpu_chunks`` and ``manifest``; other ranks return empty structures but still participate in the distributed all-gathers. + LoRA adapters are merged into their base weights here, so the export is a + plain HF checkpoint either way — see :meth:`_lora_merge_plan`. + Returns: cpu_chunks: List of ``{weight_name: cpu_tensor}`` dicts, one per shard file. manifest: Mapping of ``weight_name → shard_filename``. total_size: Total byte count of all exported tensors (for the index JSON). """ + merge_targets, adapter_keys = self._lora_merge_plan(model.model.model) + if merge_targets: + log.info( + f"[HFExportCallback] Merging {len(merge_targets)} LoRA adapter(s) into their " + "base weights; the export carries no lora_* keys." + ) + cpu_chunks: list[dict[str, torch.Tensor]] = [] manifest: dict[str, str] = {} current_chunk: dict[str, torch.Tensor] = {} current_chunk_bytes: int = 0 total_size: int = 0 file_idx: int = 0 + merged: set[str] = set() for name, param in model.model.model.named_parameters(): # Phase 2+: HFModel initialises _model via AutoModelForImageTextToText / @@ -240,12 +310,32 @@ def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], di # torch.compile and gradient-checkpointing wrappers inject prefixes into # named_parameters() output. Strip them so exported keys are HF-native, # matching what HFModel._load_vlm_weights() does for the in-memory state dict. - name = name.replace("_orig_mod.", "").replace("_checkpoint_wrapped_module.", "") + name = self._strip_wrapper_prefixes(name) + + # Adapter tensors are folded into their base weight below, so they must + # not also be written out: no HF architecture declares lora_* keys, and + # from_pretrained() drops unexpected ones with a warning — the export + # would look complete while actually being the untuned base model. + if name in adapter_keys: + continue # Gather across FSDP / TP / CP ranks (collective — all ranks must call). - if isinstance(param, torch.distributed.tensor.DTensor): - param = param.full_tensor() - param = param.detach() + param = self._gather_full(param) + + lora_module = merge_targets.get(name) + if lora_module is not None: + # lora_A / lora_B are gathered here instead of at their own + # named_parameters() entries. Every rank walks the same module + # tree in the same order, which is all the all-gather requires. + param = lora_module.merged_weight( + param, + self._gather_full(lora_module.lora_A.weight), + self._gather_full(lora_module.lora_B.weight), + ) + merged.add(name) + + # Cast after the merge: merged_weight accumulates in float32, and + # casting first would round the delta away before it is added. if self._export_dtype is not None: param = param.to(dtype=self._export_dtype) @@ -272,6 +362,37 @@ def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], di if current_chunk_bytes > 0 and is_rank0() and current_chunk: cpu_chunks.append(current_chunk) + # Every adapter the plan found must actually have been folded in. The plan + # keys off named_modules() paths and the loop off named_parameters() paths; + # if a wrapper this code does not know about ever desynchronizes the two, + # the merge silently no-ops and the export is the untuned base model. + # + # `merged` is tracked on every rank (the loop is rank-independent), so this + # aborts everywhere at the same point rather than on rank 0 alone. Both + # checks sit after the last collective, so raising cannot strand a peer + # mid-all-gather. + if len(merged) != len(merge_targets): + missed = sorted(set(merge_targets) - merged) + raise RuntimeError( + f"[HFExportCallback] LoRA merge incomplete: {len(merged)} of " + f"{len(merge_targets)} adapters folded in. Unmerged base weights: " + f"{missed[:8]}{' ...' if len(missed) > 8 else ''}. The module paths from " + "named_modules() no longer line up with the parameter paths from " + "named_parameters() — check _WRAPPER_SEGMENTS for a wrapper this code " + "does not strip." + ) + # The invariant the export must actually satisfy, asserted directly. + # manifest is rank-0-only, so this is a rank-0 check; the count above is + # what catches the desync case on every rank. + leaked = sorted(k for k in manifest if "lora_" in k) + if leaked: + raise RuntimeError( + f"[HFExportCallback] Adapter tensors leaked into the export: {leaked[:8]}" + f"{' ...' if len(leaked) > 8 else ''}. An HF checkpoint must carry merged " + "weights only; from_pretrained() would drop these and hand back the " + "untuned base model." + ) + return cpu_chunks, manifest, total_size def _save_and_upload( diff --git a/cosmos_framework/callbacks/hf_export_fsdp_test.py b/cosmos_framework/callbacks/hf_export_fsdp_test.py new file mode 100644 index 00000000..440dbd5f --- /dev/null +++ b/cosmos_framework/callbacks/hf_export_fsdp_test.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Distributed counterpart to ``hf_export_test.py`` — LoRA merging under FSDP2. + +``hf_export_test.py`` runs single-process on CPU with plain ``nn.Linear``, so it +cannot cover the part that actually carries risk: ``lora_A`` / ``lora_B`` are +DTensors sharded across ranks, and ``_gather_weights`` all-gathers them from +inside the *base weight's* loop iteration rather than at their own +``named_parameters()`` entries. That reordering is safe only because every rank +walks the module tree identically — a property worth asserting rather than +arguing. + +World size must be 2. Launch with:: + + torchrun --nproc_per_node=2 -m pytest cosmos_framework/callbacks/hf_export_fsdp_test.py + +Under plain pytest (no ``RANK``) every test skips, matching ``cfgp_ar_test`` and +``context_parallel_test``. +""" + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper +from torch.distributed.fsdp import fully_shard + +from cosmos_framework.callbacks.hf_export import HFExportCallback +from cosmos_framework.utils.generator.lora import LoraInjectedLinear + +_WORLD_SIZE = 2 + + +def setup_distributed_environment() -> tuple[int, int]: + if "RANK" not in os.environ: + pytest.skip("requires distributed environment (run with: torchrun --nproc_per_node=2)") + if not dist.is_initialized(): + dist.init_process_group(backend="nccl", init_method="env://") + rank, world_size = dist.get_rank(), dist.get_world_size() + if world_size != _WORLD_SIZE: + pytest.skip(f"requires world_size={_WORLD_SIZE}, got {world_size}") + torch.cuda.set_device(rank) + return rank, world_size + + +def _lora_linear(in_f: int, out_f: int, rank: int, alpha: int, *, bias: bool = False) -> LoraInjectedLinear: + """Materialized adapter — ``__init__`` puts lora_A / lora_B on the meta device.""" + base = nn.Linear(in_f, out_f, bias=bias) + module = LoraInjectedLinear(base, rank, alpha) + module.lora_A = nn.Linear(in_f, rank, bias=False) + module.lora_B = nn.Linear(rank, out_f, bias=False) + return module + + +class _Block(nn.Module): + """Four adapted projections plus an unadapted MLP.""" + + _ADAPTED = ("q_proj", "k_proj", "v_proj", "o_proj") + + def __init__(self, dim: int = 64, rank: int = 8, alpha: int = 16) -> None: + super().__init__() + for name in self._ADAPTED: + setattr(self, name, _lora_linear(dim, dim, rank, alpha, bias=(name == "o_proj"))) + self.mlp = nn.Linear(dim, dim * 2) + + +def _build_sharded_block() -> tuple[nn.Module, dict[str, torch.Tensor]]: + """Return an FSDP2-sharded block and the merged weights computed BEFORE sharding. + + The reference is the whole point: it is what a single-process export of the + same model would produce, so comparing against it catches any way the + sharded path could diverge. + """ + torch.manual_seed(1234) # identical init on every rank + model = _Block().cuda() + for name in _Block._ADAPTED: + module = getattr(model, name) + # lora_B ships zero-initialized; a zero adapter makes the merge a no-op + # and would let a broken merge pass. + nn.init.normal_(module.lora_A.weight, std=0.02) + nn.init.normal_(module.lora_B.weight, std=0.02) + + reference: dict[str, torch.Tensor] = {} + for name, module in model.named_modules(): + if isinstance(module, LoraInjectedLinear): + reference[f"{name}.weight"] = module.merged_weight( + module.weight.detach(), module.lora_A.weight.detach(), module.lora_B.weight.detach() + ).clone() + for name, param in model.named_parameters(): + # Adapters are folded into the base weight, so a correct export does not + # carry them and neither does the reference. + if name.endswith(("lora_A.weight", "lora_B.weight")): + continue + reference.setdefault(name, param.detach().clone()) + + # Gradient checkpointing on one projection puts a real + # _checkpoint_wrapped_module segment in the module tree — the exact shape + # that broke an earlier revision's path stripping. + model.k_proj = checkpoint_wrapper(model.k_proj) + + for child in list(model.children()): + fully_shard(child) + fully_shard(model) + return model, reference + + +def _gather(model: nn.Module) -> tuple[dict[str, torch.Tensor], dict[str, str], int]: + callback = HFExportCallback(dtype="float32") + chunks, manifest, total = callback._gather_weights(SimpleNamespace(model=SimpleNamespace(model=model))) + return {k: v for c in chunks for k, v in c.items()}, manifest, total + + +def test_adapters_are_actually_sharded(): + """Guards the test itself: without DTensors the rest proves nothing.""" + setup_distributed_environment() + model, _ = _build_sharded_block() + + sharded = [n for n, p in model.named_parameters() if isinstance(p, torch.distributed.tensor.DTensor)] + assert sharded, "FSDP2 did not shard anything; the remaining assertions would be vacuous" + assert len([n for n in sharded if "lora_" in n]) == 8, f"expected 8 sharded adapter tensors, got {sharded}" + + +def test_merged_export_matches_the_unsharded_reference(): + """The whole contract: sharded export == single-process export, key for key.""" + rank, _ = setup_distributed_environment() + model, reference = _build_sharded_block() + + flat, manifest, total = _gather(model) + # Returning on every rank is itself the assertion that the reordered + # all-gathers stay in lockstep; a mismatch hangs here instead. + dist.barrier() + + if rank != 0: + return + + assert not [k for k in flat if "lora_" in k], f"adapter keys leaked into the export: {sorted(flat)}" + assert set(flat) == set(reference), ( + f"extra={sorted(set(flat) - set(reference))} missing={sorted(set(reference) - set(flat))}" + ) + assert set(manifest) == set(flat) + assert total == sum(t.element_size() * t.numel() for t in flat.values()) + for key, tensor in flat.items(): + torch.testing.assert_close(tensor.cuda(), reference[key], rtol=0, atol=1e-5, msg=f"mismatch at {key}") + + +def test_merge_completeness_guard_fires_under_fsdp(): + """The guard must abort on every rank, not just where the manifest lives.""" + setup_distributed_environment() + model, _ = _build_sharded_block() + + callback = HFExportCallback(dtype="float32") + real_plan = callback._lora_merge_plan + # A target that no parameter name can match — i.e. the plan and the loop + # disagreeing, which is how a silently unmerged export would arise. + callback._lora_merge_plan = lambda root: ({"bogus.path.weight": None}, real_plan(root)[1]) + + with pytest.raises(RuntimeError, match="LoRA merge incomplete"): + callback._gather_weights(SimpleNamespace(model=SimpleNamespace(model=model))) diff --git a/cosmos_framework/callbacks/hf_export_test.py b/cosmos_framework/callbacks/hf_export_test.py new file mode 100644 index 00000000..eee8790d --- /dev/null +++ b/cosmos_framework/callbacks/hf_export_test.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""LoRA merging in HFExportCallback._gather_weights. + +The property under test throughout: an export taken from a LoRA run must be +indistinguishable from an export of a full fine-tune that reached the same +effective weights. Concretely — no ``lora_*`` keys, and every adapted +``.weight`` already carrying ``W + (alpha / r) * B @ A``. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from cosmos_framework.callbacks.hf_export import HFExportCallback +from cosmos_framework.utils.generator.lora import LoraInjectedLinear + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +def _lora_linear(in_features: int, out_features: int, rank: int, alpha: int, *, bias: bool = False): + """A materialized LoraInjectedLinear with a non-zero adapter. + + ``LoraInjectedLinear.__init__`` allocates lora_A / lora_B on the meta device + (production materializes them after the FSDP wrap, in + ``init_lora_weights_post_materialization``); swap in real CPU Linears here. + lora_B is deliberately non-zero — at its trained-from-zero init the merge + would be a no-op and could not distinguish a working merge from none at all. + """ + base = nn.Linear(in_features, out_features, bias=bias) + module = LoraInjectedLinear(base, rank, alpha) + module.lora_A = nn.Linear(in_features, rank, bias=False) + module.lora_B = nn.Linear(rank, out_features, bias=False) + nn.init.normal_(module.lora_A.weight, std=0.02) + nn.init.normal_(module.lora_B.weight, std=0.02) + return module + + +class _CheckpointWrapper(nn.Module): + """Mimics the attribute name gradient checkpointing injects into paths.""" + + def __init__(self, inner: nn.Module) -> None: + super().__init__() + self._checkpoint_wrapped_module = inner + + +def _fake_vlm(net: nn.Module): + """_gather_weights reaches the HF transformer at ``model.model.model``.""" + return SimpleNamespace(model=SimpleNamespace(model=net)) + + +def _gather(net: nn.Module, dtype: str = "float32") -> dict[str, torch.Tensor]: + callback = HFExportCallback(dtype=dtype) + cpu_chunks, manifest, total_size = callback._gather_weights(_fake_vlm(net)) + flat = {name: tensor for chunk in cpu_chunks for name, tensor in chunk.items()} + assert set(flat) == set(manifest), "manifest and shard contents disagree" + assert total_size == sum(t.element_size() * t.numel() for t in flat.values()) + return flat + + +# ---------------------------------------------------------------- merged_weight + + +def test_merged_weight_matches_explicit_formula(): + module = _lora_linear(8, 6, rank=4, alpha=16) + merged = module.merged_weight(module.weight, module.lora_A.weight, module.lora_B.weight) + + expected = module.weight + (16 / 4) * (module.lora_B.weight @ module.lora_A.weight) + torch.testing.assert_close(merged, expected) + + +def test_merged_weight_reproduces_the_lora_forward(): + """The merge is only correct if it is invisible to the forward pass.""" + module = _lora_linear(8, 6, rank=4, alpha=8, bias=True) + x = torch.randn(3, 8) + + merged = module.merged_weight(module.weight, module.lora_A.weight, module.lora_B.weight) + torch.testing.assert_close(F.linear(x, merged, module.bias), module(x)) + + +def test_merged_weight_is_identity_at_zero_init(): + """lora_B starts at zero, so an export before any step must be the base model.""" + module = _lora_linear(8, 6, rank=4, alpha=16) + nn.init.zeros_(module.lora_B.weight) + + merged = module.merged_weight(module.weight, module.lora_A.weight, module.lora_B.weight) + torch.testing.assert_close(merged, module.weight) + + +def test_merged_weight_preserves_base_dtype(): + module = _lora_linear(8, 6, rank=4, alpha=16) + base = module.weight.to(torch.bfloat16) + merged = module.merged_weight( + base, module.lora_A.weight.to(torch.bfloat16), module.lora_B.weight.to(torch.bfloat16) + ) + assert merged.dtype == torch.bfloat16 + + +def test_merged_weight_accumulates_in_float32(): + """A bfloat16 B @ A loses enough of the delta to be measurably worse. + + rank=64 gives the matmul enough terms for bfloat16 accumulation to drift; + the merge must land nearer the float32 result than the bfloat16 one. + """ + torch.manual_seed(0) + module = _lora_linear(64, 64, rank=64, alpha=64) + base_bf16 = module.weight.to(torch.bfloat16) + a_bf16 = module.lora_A.weight.to(torch.bfloat16) + b_bf16 = module.lora_B.weight.to(torch.bfloat16) + + merged = module.merged_weight(base_bf16, a_bf16, b_bf16).to(torch.float32) + reference = base_bf16.float() + (b_bf16.float() @ a_bf16.float()) + naive_bf16 = (base_bf16 + (b_bf16 @ a_bf16)).float() + + assert (merged - reference).abs().max() <= (naive_bf16 - reference).abs().max() + + +# ---------------------------------------------------------------- _gather_weights + + +def test_gather_weights_merges_and_drops_adapter_keys(): + net = nn.Module() + net.q_proj = _lora_linear(8, 6, rank=4, alpha=16) + net.mlp = nn.Linear(6, 6) + expected = net.q_proj.merged_weight(net.q_proj.weight, net.q_proj.lora_A.weight, net.q_proj.lora_B.weight) + + flat = _gather(net) + + assert not [k for k in flat if "lora_" in k], f"adapter keys leaked into the export: {sorted(flat)}" + assert set(flat) == {"q_proj.weight", "mlp.weight", "mlp.bias"} + torch.testing.assert_close(flat["q_proj.weight"], expected) + + +def test_gather_weights_exports_lora_bias_unchanged(): + net = nn.Module() + net.q_proj = _lora_linear(8, 6, rank=4, alpha=16, bias=True) + + flat = _gather(net) + + assert set(flat) == {"q_proj.weight", "q_proj.bias"} + torch.testing.assert_close(flat["q_proj.bias"], net.q_proj.bias) + + +def test_gather_weights_merges_under_a_checkpoint_wrapper(): + """Adapters under a gradient-checkpointing wrapper must still be found. + + The module path and the parameter path both carry the wrapper segment; the + merge only fires if the two are stripped consistently. + """ + lora = _lora_linear(8, 6, rank=4, alpha=16) + net = nn.Module() + net.layer = _CheckpointWrapper(lora) + expected = lora.merged_weight(lora.weight, lora.lora_A.weight, lora.lora_B.weight) + + flat = _gather(net) + + assert set(flat) == {"layer.weight"} + torch.testing.assert_close(flat["layer.weight"], expected) + + +def test_gather_weights_leaves_a_full_finetune_untouched(): + """No LoraInjectedLinear anywhere means the pre-existing path is unchanged.""" + net = nn.Sequential(nn.Linear(8, 6), nn.Linear(6, 4)) + reference = {name: param.detach().clone() for name, param in net.named_parameters()} + + flat = _gather(net) + + assert set(flat) == set(reference) + for name, tensor in flat.items(): + torch.testing.assert_close(tensor, reference[name]) + + +def test_gather_weights_raises_when_a_wrapper_desyncs_the_paths(): + """The guard against the failure mode this code is most exposed to. + + _lora_merge_plan keys off named_modules() paths and the loop off + named_parameters() paths. Simulate an unknown wrapper by stripping a segment + the plan does not strip: the merge would silently no-op and export the base + model, so it must abort instead. + """ + lora = _lora_linear(8, 6, rank=4, alpha=16) + net = nn.Module() + net.layer = _CheckpointWrapper(lora) + + callback = HFExportCallback(dtype="float32") + # A wrapper the plan does not strip leaves its segment in the plan's keys + # while the loop's parameter names have it stripped — the two never meet. + # (This is the real bug an earlier revision of this code shipped.) + wrapped = "layer._checkpoint_wrapped_module" + callback._lora_merge_plan = lambda root: ( + {f"{wrapped}.weight": lora}, + {f"{wrapped}.lora_A.weight", f"{wrapped}.lora_B.weight"}, + ) + + # The count check fires before the leak check, naming the actual cause. + with pytest.raises(RuntimeError, match="LoRA merge incomplete"): + callback._gather_weights(_fake_vlm(net)) + + +def test_gather_weights_raises_when_adapter_keys_would_leak(): + """Belt-and-braces on the invariant itself, independent of the count check.""" + net = nn.Module() + net.q_proj = _lora_linear(8, 6, rank=4, alpha=16) + + callback = HFExportCallback(dtype="float32") + real_plan = callback._lora_merge_plan + # Keep merge_targets intact (so the count check passes) but forget to exclude + # the adapter parameters. + callback._lora_merge_plan = lambda root: (real_plan(root)[0], set()) + + with pytest.raises(RuntimeError, match="Adapter tensors leaked"): + callback._gather_weights(_fake_vlm(net)) + + +def test_gather_weights_casts_to_the_export_dtype(): + net = nn.Module() + net.q_proj = _lora_linear(8, 6, rank=4, alpha=16) + + flat = _gather(net, dtype="bfloat16") + + assert all(t.dtype == torch.bfloat16 for t in flat.values()) diff --git a/cosmos_framework/configs/base/reasoner/defaults/policy_config.py b/cosmos_framework/configs/base/reasoner/defaults/policy_config.py index 8c5adccf..3989f1cb 100644 --- a/cosmos_framework/configs/base/reasoner/defaults/policy_config.py +++ b/cosmos_framework/configs/base/reasoner/defaults/policy_config.py @@ -30,8 +30,30 @@ class PolicyConfig: # 0 < exponent < 1 -> interpolation; e.g. exponent=0.5 gives square-root per-token loss (Qwen3-VL) weighted_ce_exponent: float = 1.0 + # LoRA (parameter-efficient fine-tuning). When ``lora_enabled=True``, + # ``VLMModel._init_vlm`` injects the custom LoRA adapters BEFORE FSDP wrap on + # the meta-device HF backbone, then re-initializes lora_A/lora_B after the + # meta tensors are materialized and the base weights are loaded. Pair with + # ``optimizer.keys_to_select=["lora_"]``. + # + # ``lora_target_modules`` is matched by EXACT child name (see + # ``_inject_lora_inplace``), not substring. The default targets the four + # Qwen3-VL LLM attention projections; the vision tower names its projections + # ``qkv`` / ``proj`` / ``linear_fc1`` / ``linear_fc2``, so the ViT is never + # touched. Other model families (e.g. cosmos3_edge) may need different names. + # + # ``lora_exclude_path_regex`` is searched against each candidate module's + # dotted path and skips matches. Needed when the two towers of a VLM share + # projection names: Cosmos3-Edge names its LLM projections q/k/v/o_proj and + # its SigLIP2 vision projections q/k/v/out_proj, so three of four collide and + # name matching alone would inject adapters into the (frozen) vision tower. + lora_enabled: bool = False + lora_rank: int = 16 + lora_alpha: int = 32 + lora_target_modules: str = "q_proj,k_proj,v_proj,o_proj" + lora_exclude_path_regex: str = "" + # Extra model config - lora: Union[str, None] = None enable_liger_kernel: bool = False trainable_map: Union[str, None] = None monkey_patch_for_text_only_data: bool = False diff --git a/cosmos_framework/configs/toml_config/sft_config.py b/cosmos_framework/configs/toml_config/sft_config.py index 04d0efbc..ee55f479 100644 --- a/cosmos_framework/configs/toml_config/sft_config.py +++ b/cosmos_framework/configs/toml_config/sft_config.py @@ -346,12 +346,13 @@ class ModelConfig(BaseModel): lora_enabled: bool = Field( default=False, description=( - "Inject LoRA adapters into the generation pathway BEFORE FSDP " - "wraps the network. Pair with optimizer.keys_to_select=['lora_'] " - "(train only adapters) and checkpoint.keys_to_skip_loading=[" - "..., 'lora_'] (don't load missing adapter tensors). Used by " - "SUPER-tier configs (e.g. vision_sft_super); NANO-tier leaves " - "it off. Skipped on VLM." + "Inject LoRA adapters BEFORE FSDP wraps the network. Pair with " + "optimizer.keys_to_select=['lora_'] (train only adapters) and " + "checkpoint.keys_to_skip_loading=[..., 'lora_'] (don't load " + "missing adapter tensors). On VFM this targets the generation " + "pathway (e.g. vision_sft_super); on VLM it targets the HF " + "backbone via model.config.policy.lora_* (e.g. " + "examples/toml/sft_config/videophy2_lora_super.toml)." ), ) lora_rank: int = Field( @@ -371,8 +372,22 @@ class ModelConfig(BaseModel): lora_target_modules: str = Field( default="q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen", description=( - "Comma-separated substrings of param names that get a LoRA " - "adapter. Defaults target the four MoE-gen projection matrices." + "Comma-separated EXACT child-module names that get a LoRA " + "adapter (matched by name, not substring). The default targets " + "the four MoE-gen projection matrices, which is the VFM layout; " + "VLM recipes override it (Qwen3-VL: " + "'q_proj,k_proj,v_proj,o_proj')." + ), + ) + lora_exclude_path_regex: str = Field( + default="", + description=( + "Regex searched against each candidate module's dotted path; " + "matches are skipped. Use when the two towers of a VLM share " + "projection names — Cosmos3-Edge names its LLM projections " + "q/k/v/o_proj and its SigLIP2 vision projections q/k/v/out_proj, " + "so name matching alone would also adapt the (frozen) vision " + "tower. Empty = no exclusion. **VLM only.**" ), ) diff --git a/cosmos_framework/configs/toml_config/toml_config_helper.py b/cosmos_framework/configs/toml_config/toml_config_helper.py index 4d1535c5..e03bee84 100644 --- a/cosmos_framework/configs/toml_config/toml_config_helper.py +++ b/cosmos_framework/configs/toml_config/toml_config_helper.py @@ -16,6 +16,7 @@ from __future__ import annotations +import re from typing import Any @@ -52,6 +53,7 @@ ("job", "upload_reproducible_setup"): ("upload_reproducible_setup",), ("model", "attn_implementation"): None, ("model", "backbone"): None, # VLM-only — VFM has no model.config.backbone + ("model", "lora_exclude_path_regex"): None, # VLM-only — VFM's single tower needs no path scoping # Per-caption token cap lives on the nested SFT dataset, not a top-level # dataloader scalar — route it to the get_sft_dataset node. ("dataloader_train", "max_caption_tokens"): ( @@ -72,10 +74,6 @@ # No VLM analog — skip these leaves ("model", "max_num_tokens_after_packing"): None, ("model", "joint_attn_implementation"): None, - ("model", "lora_enabled"): None, - ("model", "lora_rank"): None, - ("model", "lora_alpha"): None, - ("model", "lora_target_modules"): None, ("model", "tokenizer"): None, # blocks model.tokenizer.* ("dataloader_train", "seed"): None, ("optimizer", "eps"): None, # VLM_OPTIMIZER_KWARGS has no eps field @@ -85,6 +83,13 @@ ("model", "attn_implementation"): ("model", "config", "policy", "attn_implementation"), ("model", "ema"): ("model", "config", "ema"), ("model", "backbone"): ("model", "config", "policy", "backbone"), + # LoRA knobs live on PolicyConfig for VLM (they sit flat on + # OmniMoTModelConfig for VFM, which the catch-all below already handles). + ("model", "lora_enabled"): ("model", "config", "policy", "lora_enabled"), + ("model", "lora_rank"): ("model", "config", "policy", "lora_rank"), + ("model", "lora_alpha"): ("model", "config", "policy", "lora_alpha"), + ("model", "lora_target_modules"): ("model", "config", "policy", "lora_target_modules"), + ("model", "lora_exclude_path_regex"): ("model", "config", "policy", "lora_exclude_path_regex"), # VLM uses CosmosDataLoader whose batch/token caps live on the nested # PoolPackingBatcher (dataloader_train.batcher.*), not flat on the loader. ("dataloader_train", "max_samples_per_batch"): ("dataloader_train", "batcher", "max_batch_size"), @@ -181,6 +186,11 @@ def _emit_with_remap( out.append(f"{'.'.join(new_path)}={_hydra_format(value)}") +# Characters Hydra accepts in an unquoted override value. Anything else (regex +# metacharacters, commas, whitespace, quotes) gets single-quoted by _hydra_format. +_HYDRA_UNQUOTED_SAFE = re.compile(r"[A-Za-z0-9_./:@+-]*") + + def _hydra_format(v: Any, in_list: bool = False) -> str: """Convert a Python value to a Hydra CLI override RHS. @@ -199,12 +209,17 @@ def _hydra_format(v: Any, in_list: bool = False) -> str: return "[" + ",".join(_hydra_format(x, in_list=True) for x in v) + "]" if isinstance(v, str): # Inside a list literal, always quote so numeric-looking strings - # ("480") aren't parsed as int. At top level, quote only when the - # string contains characters Hydra would otherwise interpret — - # commas (sweep / list marker) or whitespace. Env-interpolation - # strings like ``${oc.env:NAME}`` are safe unquoted because Hydra - # recognizes the ``${...}`` form even with a colon inside. - if in_list or "," in v or " " in v: + # ("480") aren't parsed as int. At top level, quote anything outside + # Hydra's unquoted-value character set: a comma is a sweep/list marker, + # whitespace ends the token, and metacharacters like ``^`` or ``\`` + # (regex values such as lora_exclude_path_regex) make its lexer throw + # LexerNoViableAltException. Single quotes preserve backslashes + # verbatim, so a quoted regex round-trips unchanged. + # Env-interpolation strings like ``${oc.env:NAME}`` must stay unquoted + # so Hydra resolves them instead of passing the literal through. + if not in_list and v.startswith("${") and v.endswith("}"): + return v + if in_list or not _HYDRA_UNQUOTED_SAFE.fullmatch(v): return f"'{v}'" return v return str(v) diff --git a/cosmos_framework/model/generator/vlm_model.py b/cosmos_framework/model/generator/vlm_model.py index 2d33fa4f..c0cecb99 100644 --- a/cosmos_framework/model/generator/vlm_model.py +++ b/cosmos_framework/model/generator/vlm_model.py @@ -390,6 +390,23 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: if policy.backbone.pretrained_weights.backbone_path: _get_overlay_config(hf_model.hf_config.model_type) + # ── b.2. Inject LoRA adapters (still on meta, still pre-FSDP) ── + # Ordering is load-bearing: the injector must see UNSHARDED nn.Linear + # shapes. Injecting after ``parallelize()`` builds ``lora_B`` at the + # per-rank shard size (e.g. 8192/4=2048) and crashes at forward time. + # ``lora_A``/``lora_B`` stay uninitialized on meta here; step g.3 + # initializes them once the tensors are real. + if policy.lora_enabled: + from cosmos_framework.utils.generator.lora import inject_lora_pre_fsdp + + inject_lora_pre_fsdp( + hf_model.model, + lora_rank=policy.lora_rank, + lora_alpha=policy.lora_alpha, + lora_target_modules=policy.lora_target_modules, + lora_exclude_path_regex=policy.lora_exclude_path_regex or None, + ) + # ── c. Build ParallelDims + device mesh ── # Overlay-mesh design (see vfm/utils/parallelism.py): cp/cfgp do NOT # consume FSDP rank slots, so dp_replicate * dp_shard == world_size @@ -468,6 +485,13 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: hf_model.tie_embeddings() # ── g. Load pretrain weights ── + # LoRA adapters exist on the model but never in a pretrained checkpoint. + # ``load_vlm_model``'s Phase-6 completeness check raises on any model key + # the checkpoint lacks, so the adapter keys must be tolerated explicitly. + # Patterns are applied with ``re.fullmatch`` against the resolved model + # key, hence the leading ``.*``. + lora_skip_patterns = [r".*\.lora_[AB]\.weight"] if policy.lora_enabled else [] + if load_pretrain_weights: if policy.backbone.safetensors_path: safetensors_local_path = maybe_download_hf_model_from_s3( @@ -483,6 +507,7 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: checkpoint_path=safetensors_local_path, credential_path=None, # local path after download parallel_dims=parallel_dims if torch.distributed.is_initialized() else None, + extra_skip_patterns=lora_skip_patterns or None, ) # ── g.2. Optional LLM overlay (backbone.pretrained_weights) ── @@ -508,7 +533,7 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: checkpoint_path=llm_local_path, credential_path=None, parallel_dims=parallel_dims if torch.distributed.is_initialized() else None, - extra_skip_patterns=overlay_skip_patterns, + extra_skip_patterns=overlay_skip_patterns + lora_skip_patterns, ) lm_loaded = {k for k in keys_loaded if is_lm_key(k)} if not lm_loaded: @@ -520,6 +545,21 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: ) log.info(f"VLMModel: overlaid {len(lm_loaded)} language-model params from {llm_path}") + # ── g.3. Initialize LoRA adapters ── + # Must run AFTER step e (meta -> real CUDA storage) and after every + # load_weights (which skips these keys, leaving whatever ``empty_like`` + # allocated). ``lora_A ~ kaiming_uniform_``, ``lora_B = 0`` — so the + # adapter contributes exactly zero on the first forward and the run starts + # from the pretrained model's loss. + # + # Ordering vs step h: the Parakeet load below writes only into + # ``sound_und_model.encoder`` and its checkpoint has no ``lora_*`` keys, + # so it cannot clobber what this initializes. + if policy.lora_enabled: + from cosmos_framework.utils.generator.lora import init_lora_weights_post_materialization + + init_lora_weights_post_materialization(hf_model.model) + # ── h. Load the immutable standalone Parakeet artifact ── # This runs for both fresh starts and DCP resumes. On resume, DCP may # subsequently restore the same encoder when it was checkpointed; when diff --git a/cosmos_framework/utils/generator/lora.py b/cosmos_framework/utils/generator/lora.py index bbd3215d..5ae8a485 100644 --- a/cosmos_framework/utils/generator/lora.py +++ b/cosmos_framework/utils/generator/lora.py @@ -14,6 +14,7 @@ from __future__ import annotations import math +import re import torch import torch.nn as nn @@ -65,6 +66,23 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: lora_out = self.lora_B(self.lora_A(x)) return base_out + self._lora_scale * lora_out + def merged_weight(self, base: torch.Tensor, lora_a: torch.Tensor, lora_b: torch.Tensor) -> torch.Tensor: + """Fold the adapter into the base weight: ``W + (alpha / r) * B @ A``. + + The three tensors are passed in rather than read off ``self`` because the + only caller (``HFExportCallback``) has already all-gathered them out of + their FSDP shards — ``self.weight`` is still a per-rank ``DTensor`` at + that point. Only ``_lora_scale`` comes from the module, so the merge + stays in lockstep with :meth:`forward` if the scaling convention changes. + + The accumulation runs in float32 even when the export dtype is bfloat16: + the delta is typically orders of magnitude smaller than the base weight, + so adding it in bfloat16 rounds much of it away. The result is cast back + to ``base``'s dtype. + """ + delta = torch.mm(lora_b.to(torch.float32), lora_a.to(torch.float32)) + return (base.to(torch.float32) + self._lora_scale * delta).to(base.dtype) + def _target_matches(full_child_path: str, child_name: str, target: str) -> bool: """Return True if ``target`` selects the child at ``full_child_path``. @@ -91,6 +109,7 @@ def _inject_lora_inplace( target_modules: list[str], rank: int, alpha: int, + exclude_path_regex: str | None = None, ) -> int: """Replace each targeted ``nn.Linear`` child in-place with ``LoraInjectedLinear``. @@ -101,16 +120,32 @@ def _inject_lora_inplace( Snapshots ``named_modules()`` before mutating the tree so newly-inserted LoRA submodules are not re-visited. + + ``exclude_path_regex`` skips any module whose dotted path matches (searched, + not fullmatch). Name matching alone cannot always separate two towers of a + VLM: Cosmos3-Edge names its LLM projections ``q_proj``/``k_proj``/``v_proj``/ + ``o_proj`` and its SigLIP2 vision projections ``q_proj``/``k_proj``/ + ``v_proj``/``out_proj`` — three of the four names collide. Passing + ``r"^model\\.visual\\."`` keeps the adapters out of the vision tower. """ + exclude = re.compile(exclude_path_regex) if exclude_path_regex else None replaced = 0 for parent_name, parent in list(network.named_modules()): for child_name, child in list(parent.named_children()): if not isinstance(child, nn.Linear): continue full_child_path = f"{parent_name}.{child_name}" if parent_name else child_name - if any(_target_matches(full_child_path, child_name, t) for t in target_modules): - setattr(parent, child_name, LoraInjectedLinear(child, rank, alpha)) - replaced += 1 + if not any(_target_matches(full_child_path, child_name, t) for t in target_modules): + continue + # Selection and exclusion are orthogonal, and exclusion runs second: + # `_target_matches` picks by name or path suffix, then the regex + # carves a subtree back out. Cosmos3-Edge needs both — its LLM and + # SigLIP2 tower share three of four projection names, so no target + # spelling separates them, but the tower is one contiguous subtree. + if exclude is not None and exclude.search(full_child_path): + continue + setattr(parent, child_name, LoraInjectedLinear(child, rank, alpha)) + replaced += 1 return replaced @@ -120,6 +155,7 @@ def inject_lora_pre_fsdp( lora_rank: int, lora_alpha: int, lora_target_modules: str, + lora_exclude_path_regex: str | None = None, ) -> torch.nn.Module: """Inject LoRA adapters into ``network`` BEFORE FSDP wrap on meta device. @@ -161,10 +197,15 @@ def _target_exists(t: str) -> bool: if invalid_modules: log.warning(f"LoRA target modules not found in model: {invalid_modules}") - log.info(f"Injecting LoRA on meta device: rank={lora_rank}, alpha={lora_alpha}, targets={target_modules_list}") + log.info( + f"Injecting LoRA on meta device: rank={lora_rank}, alpha={lora_alpha}, " + f"targets={target_modules_list}, exclude_path_regex={lora_exclude_path_regex!r}" + ) try: - replaced = _inject_lora_inplace(network, target_modules_list, lora_rank, lora_alpha) + replaced = _inject_lora_inplace( + network, target_modules_list, lora_rank, lora_alpha, exclude_path_regex=lora_exclude_path_regex + ) except Exception as e: raise RuntimeError(f"Failed to inject LoRA adapters into model: {e}") from e diff --git a/docs/faq.md b/docs/faq.md index 445bab50..e1787d99 100644 --- a/docs/faq.md +++ b/docs/faq.md @@ -268,7 +268,7 @@ Knobs are in the recipe TOML under `[model]`, `[model.parallelism]`, and `[datal 5. **Lower `[dataloader_train].max_samples_per_batch`** to cap samples per micro-batch. `None` lets the packer's token budget decide; setting an explicit small number trades throughput for headroom. -6. **Enable LoRA on a Cosmos3-Nano recipe.** Nano recipes are full-finetune by default (`lora_enabled = false`); setting `[model].lora_enabled = true` trains low-rank adapters instead of the full weights, dropping optimizer-state memory substantially. The `_super` recipes (e.g. `vision_sft_super`) are already LoRA-only, so this lever doesn't apply there. +6. **Enable LoRA.** Setting `[model].lora_enabled = true` trains low-rank adapters instead of the full weights, dropping optimizer-state memory substantially. Generator (VFM) nano recipes are full-finetune by default, so this lever applies there; the `_super` generator recipes (e.g. `vision_sft_super`) are already LoRA-only. It also works on reasoner (VLM) recipes — see `examples/toml/sft_config/videophy2_lora_{super,edge}.toml`. Pair it with `[optimizer].keys_to_select = ["lora_"]`, and on a VLM whose vision tower shares projection names with its LLM (Cosmos3-Edge), add `[model].lora_exclude_path_regex` so the frozen tower is not adapted too. See [docs/training.md](./training.md) for the full SFT setup and TOML reference (`[model.activation_checkpointing]`, `[model.parallelism]`, `[dataloader_train]` sections). diff --git a/docs/sft_config.md b/docs/sft_config.md index 241f90fe..76b81bae 100644 --- a/docs/sft_config.md +++ b/docs/sft_config.md @@ -84,16 +84,17 @@ Run identity + meta-fields that pick the Hydra config tree to load. Top-level model knobs. Lands at `model.config.*` on VFM and on VLM; sub-tree paths are remapped per the [VFM ↔ VLM path remaps](#vfm--vlm-path-remaps). -| field | default | description | -| ------------------------------ | --------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -| `max_num_tokens_after_packing` | `13312` | Token-packing target: max tokens after sequence packing. `-1` disables the cap. **VFM only** — VLM uses `data_setting.max_tokens` (tail override). | -| `joint_attn_implementation` | `"two_way"` | VFM attention layout: `"two_way"` (U/G blocks with cross-attention), `"three_way"` (adds sparsity-aware third block — NATTEN), or `"flex"` (legacy). **VFM only.** | -| `attn_implementation` | `"cosmos"` | VLM HF attention impl: `"cosmos"` (NATTEN/Blackwell-FMHA wrapper), `"flash_attention_2"`, `"sdpa"`, or `"eager"`. **VLM only.** | -| `lora_enabled` | `false` | Inject LoRA adapters into the generation pathway BEFORE FSDP wraps the network. Pair with `optimizer.keys_to_select=["lora_"]` and `checkpoint.keys_to_skip_loading=[…, "lora_"]`. Used by SUPER-tier recipes; NANO leaves it off. **VFM only.** | -| `lora_rank` | `16` | LoRA rank `r`. Adapter shape is (rank × hidden_dim) per target module. Typical: 4 / 8 / 16 / 32. | -| `lora_alpha` | `32` | LoRA scaling factor. Effective magnitude is `alpha / rank`; rank=16 alpha=32 → 2× scale. | -| `lora_target_modules` | `"q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen"` | Comma-separated substrings of param names that receive an adapter. Default targets the four MoE-gen projection matrices. | -| `precision` | `"bfloat16"` | Compute dtype for forward/backward (`MixedPrecisionPolicy.param_dtype`). `"bfloat16"` is standard for Hopper/Blackwell. (Was `[model.parallelism].precision` before the `ParallelismConfig` split.) | +| field | default | description | +| ------------------------------ | --------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `max_num_tokens_after_packing` | `13312` | Token-packing target: max tokens after sequence packing. `-1` disables the cap. **VFM only** — VLM uses `data_setting.max_tokens` (tail override). | +| `joint_attn_implementation` | `"two_way"` | VFM attention layout: `"two_way"` (U/G blocks with cross-attention), `"three_way"` (adds sparsity-aware third block — NATTEN), or `"flex"` (legacy). **VFM only.** | +| `attn_implementation` | `"cosmos"` | VLM HF attention impl: `"cosmos"` (NATTEN/Blackwell-FMHA wrapper), `"flash_attention_2"`, `"sdpa"`, or `"eager"`. **VLM only.** | +| `lora_enabled` | `false` | Inject LoRA adapters BEFORE FSDP wraps the network. Pair with `optimizer.keys_to_select=["lora_"]` and `checkpoint.keys_to_skip_loading=[…, "lora_"]`. On **VFM** this targets the generation pathway (e.g. `vision_sft_super`); on **VLM** it targets the HF reasoner backbone via `model.config.policy.lora_*` (e.g. `videophy2_lora_super`). NANO-tier generator recipes leave it off. | +| `lora_rank` | `16` | LoRA rank `r`. Adapter shape is (rank × hidden_dim) per target module. Typical: 4 / 8 / 16 / 32. | +| `lora_alpha` | `32` | LoRA scaling factor. Effective magnitude is `alpha / rank`; rank=16 alpha=32 → 2× scale. | +| `lora_target_modules` | `"q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen"` | Comma-separated module selectors, matched **not** by substring: a plain name (`q_proj_moe_gen`) matches a module whose own child name is exactly that; a name containing a `.` (`mlp_moe_gen.up_proj`) matches by full-path suffix, which disambiguates leaf names shared across towers. The default targets the four MoE-gen projections (VFM); VLM recipes use `"q_proj,k_proj,v_proj,o_proj"`. | +| `lora_exclude_path_regex` | `""` (disabled) | Skip any module whose dotted path matches this regex (searched, not fullmatched), applied **after** target matching. Needed when name matching alone cannot separate two towers: Cosmos3-Edge names its LLM projections `q_proj`/`k_proj`/`v_proj`/`o_proj` and its SigLIP2 vision projections `q_proj`/`k_proj`/`v_proj`/`out_proj` — three of four collide, so `"^model\\.visual\\."` is what keeps the adapters out of the frozen vision tower. **VLM only** (VFM's single tower needs no path scoping). | +| `precision` | `"bfloat16"` | Compute dtype for forward/backward (`MixedPrecisionPolicy.param_dtype`). `"bfloat16"` is standard for Hopper/Blackwell. (Was `[model.parallelism].precision` before the `ParallelismConfig` split.) | ### `[model.ema]` @@ -279,23 +280,25 @@ vae_path = "${oc.env:WAN_VAE_PATH}" The same TOML key lands at different Hydra paths depending on `[job].task`: -| TOML path | VFM (`task="vfm"`) Hydra path | VLM (`task="vlm"`) Hydra path | -| ---------------------------------------------------------------------------------------- | ----------------------------------------- | ----------------------------------------- | -| `job.upload_reproducible_setup` | `upload_reproducible_setup` | `upload_reproducible_setup` | -| `model.` | `model.config.` | `model.config.` | -| `model.parallelism.*` | `model.config.parallelism.*` | `model.config.parallelism.*` | -| `model.compile.*` | `model.config.compile.*` | `model.config.compile.*` | -| `model.activation_checkpointing.*` | `model.config.activation_checkpointing.*` | `model.config.activation_checkpointing.*` | -| `model.precision` | `model.config.precision` | `model.config.precision` | -| `model.attn_implementation` | *(skipped — VLM-only)* | `model.config.policy.attn_implementation` | -| `model.backbone.*` | *(skipped — VLM-only)* | `model.config.policy.backbone.*` | -| `model.ema.*` | `model.config.ema.*` | `model.config.ema.*` | -| `model.tokenizer.*` | `model.config.tokenizer.*` | *(skipped — VFM-only)* | -| `model.{max_num_tokens_after_packing, joint_attn_implementation, lora_*}` | passes through | *(skipped — VFM-only)* | -| `dataloader_train.max_samples_per_batch` | passes through | `dataloader_train.max_batch_size` | -| `dataloader_train.max_sequence_length` | passes through | `dataloader_train.max_tokens` | -| `dataloader_train.seed` | passes through | *(skipped — VLM has no seed kwarg)* | -| `optimizer.eps`, `scheduler.verbosity_interval`, `trainer.callbacks.compile_tokenizer.*` | passes through | *(skipped — VLM has no analog)* | +| TOML path | VFM (`task="vfm"`) Hydra path | VLM (`task="vlm"`) Hydra path | +| ---------------------------------------------------------------------------------------- | ----------------------------------------- | --------------------------------------------- | +| `job.upload_reproducible_setup` | `upload_reproducible_setup` | `upload_reproducible_setup` | +| `model.` | `model.config.` | `model.config.` | +| `model.parallelism.*` | `model.config.parallelism.*` | `model.config.parallelism.*` | +| `model.compile.*` | `model.config.compile.*` | `model.config.compile.*` | +| `model.activation_checkpointing.*` | `model.config.activation_checkpointing.*` | `model.config.activation_checkpointing.*` | +| `model.precision` | `model.config.precision` | `model.config.precision` | +| `model.attn_implementation` | *(skipped — VLM-only)* | `model.config.policy.attn_implementation` | +| `model.backbone.*` | *(skipped — VLM-only)* | `model.config.policy.backbone.*` | +| `model.ema.*` | `model.config.ema.*` | `model.config.ema.*` | +| `model.tokenizer.*` | `model.config.tokenizer.*` | *(skipped — VFM-only)* | +| `model.{max_num_tokens_after_packing, joint_attn_implementation}` | passes through | *(skipped — VFM-only)* | +| `model.lora_{enabled,rank,alpha,target_modules}` | passes through | `model.config.policy.lora_*` | +| `model.lora_exclude_path_regex` | *(skipped — VLM-only)* | `model.config.policy.lora_exclude_path_regex` | +| `dataloader_train.max_samples_per_batch` | passes through | `dataloader_train.max_batch_size` | +| `dataloader_train.max_sequence_length` | passes through | `dataloader_train.max_tokens` | +| `dataloader_train.seed` | passes through | *(skipped — VLM has no seed kwarg)* | +| `optimizer.eps`, `scheduler.verbosity_interval`, `trainer.callbacks.compile_tokenizer.*` | passes through | *(skipped — VLM has no analog)* | Authoritative source: `PATH_REMAPS` in [`toml_config_helper.py`](../cosmos_framework/configs/toml_config/toml_config_helper.py). diff --git a/docs/training.md b/docs/training.md index 854a6fce..90cb9a6a 100644 --- a/docs/training.md +++ b/docs/training.md @@ -166,6 +166,37 @@ python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ +
Reasoner Alignment SFT with VideoPhy-2, LoRA (Cosmos3-Super / Cosmos3-Edge) + +LoRA counterparts of the VideoPhy-2 recipes above: same dataset and dataflow, but the reasoner backbone is +frozen and only rank-16 adapters on the LLM attention projections train +(`optimizer.keys_to_select = ["lora_"]`). Optimizer state is adapter-sized rather than +backbone-sized, which is what lets the 32B Super tier sit comfortably on a 4-GPU allocation. + +Both select the **full-fine-tune** experiment (`[job].experiment = "videophy2_sft_{super,edge}"`) and switch +LoRA on through TOML overrides — there is no separate LoRA experiment to register. + +| Launch shell | Tier | Notes | +| --------------------------------------------- | -------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `examples/launch_sft_videophy2_lora_super.sh` | Qwen3-VL-32B | Public snapshot resolved from `model_name`; no converter step. | +| `examples/launch_sft_videophy2_lora_edge.sh` | Nemotron-2B-Dense-VL | Needs `lora_exclude_path_regex = "^model\\.visual\\."` — its LLM and SigLIP2 tower share three of four projection names, so name matching alone would adapt the frozen vision tower too. | + +For the 8B tier, point `[job].experiment` at `videophy2_sft_nano` and `[model.backbone].model_name` at +`Qwen/Qwen3-VL-8B-Instruct` in the super TOML; nothing else changes. + +Checkpoints: the launch shells pass `checkpoint.hf_export.enabled=false` because these are convergence runs +and a full HF snapshot per save is wasted work. When you do want one, drop that override — `HFExportCallback` +merges the adapter into the base weights (`W + (alpha/r) · B·A`), so a LoRA export is an ordinary HF +checkpoint with no `lora_*` keys, loadable exactly like a full fine-tune's. + +```shell +# Step 1 (data): same as the non-LoRA recipes. +python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ + --out_root examples/data/videophysics --split both +``` + +
+
Action-Policy Post-Training (DROID / LIBERO) Robot action-policy recipes: DROID (`joint_pos` 8-D actions + proprioceptive state) and @@ -235,15 +266,17 @@ bash examples/launch_sft_vision_nano.sh Each launcher's default paths come from the `DATASET_PATH` + `BASE_CHECKPOINT_PATH` defaults declared at the top of its `.sh` (each uses `: "${VAR:=…}"` so any value you `export` in the shell before launching wins over the default): -| Launch shell | Post-Training Task | Default $DATASET_PATH (under examples/data/) | Default $BASE_CHECKPOINT_PATH (under examples/checkpoints/) | -| ------------------------------ | ------------------ | ---------------------------------------------------------- | ------------------------------------------------------------------ | -| `launch_sft_vision_nano.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Nano` | -| `launch_sft_vision_super.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Super` | -| `launch_sft_vision_edge.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Edge` | -| `launch_sft_llava_ov.sh` | Reasoner SFT | (none; dataset streams from HF Hub) | (none; backbone fetched at startup, or set `VLM_SAFETENSORS_PATH`) | -| `launch_sft_videophy2_nano.sh` | Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; set `VLM_SAFETENSORS_PATH` env) | -| `launch_sft_videophy2_super.sh`| Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; set `VLM_SAFETENSORS_PATH` env — Cosmos3-Super-VLM merge) | -| `launch_sft_videophy2_edge.sh` | Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; weights load directly from `nvidia/Cosmos3-Edge`) | +| Launch shell | Post-Training Task | Default $DATASET_PATH (under examples/data/) | Default $BASE_CHECKPOINT_PATH (under examples/checkpoints/) | +| ------------------------------------ | ------------------- | ---------------------------------------------------------- | -------------------------------------------------------------------------- | +| `launch_sft_vision_nano.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Nano` | +| `launch_sft_vision_super.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Super` | +| `launch_sft_vision_edge.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Edge` | +| `launch_sft_llava_ov.sh` | Reasoner SFT | (none; dataset streams from HF Hub) | (none; backbone fetched at startup, or set `VLM_SAFETENSORS_PATH`) | +| `launch_sft_videophy2_nano.sh` | Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; set `VLM_SAFETENSORS_PATH` env) | +| `launch_sft_videophy2_super.sh` | Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; set `VLM_SAFETENSORS_PATH` env — Cosmos3-Super-VLM merge) | +| `launch_sft_videophy2_edge.sh` | Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; weights load directly from `nvidia/Cosmos3-Edge`) | +| `launch_sft_videophy2_lora_super.sh` | Reasoner SFT (LoRA) | (none; set `VIDEOPHYSICS_ROOT` env) | (none; public `Qwen/Qwen3-VL-32B-Instruct`, or set `VLM_SAFETENSORS_PATH`) | +| `launch_sft_videophy2_lora_edge.sh` | Reasoner SFT (LoRA) | (none; set `VIDEOPHYSICS_ROOT` env) | (none; weights load directly from `nvidia/Cosmos3-Edge`) | `WAN_VAE_PATH` defaults to `examples/checkpoints/wan22_vae/Wan2.2_VAE.pth` for every non-reasoner recipe. @@ -424,7 +457,8 @@ The commonly tuned knobs: 1. `max_num_tokens_after_packing` — VFM token-packing target. `-1` disables the cap. VFM only; VLM uses `data_setting.max_tokens` (tail override). 1. `joint_attn_implementation` — VFM attention layout: `"two_way"` / `"three_way"` (NATTEN) / `"flex"`. 1. `attn_implementation` — VLM attention impl: `"cosmos"` / `"flash_attention_2"` / `"sdpa"` / `"eager"`. VLM only. - 1. `lora_enabled`, `lora_rank`, `lora_alpha`, `lora_target_modules` — LoRA adapter knobs for the generation pathway. Used by SUPER-tier recipes; NANO-tier leaves `lora_enabled=false`. VFM only. + 1. `lora_enabled`, `lora_rank`, `lora_alpha`, `lora_target_modules` — LoRA adapter knobs. On VFM they target the generation pathway (SUPER-tier recipes use them; NANO-tier leaves `lora_enabled=false`); on VLM they target the HF reasoner backbone (`videophy2_lora_{super,edge}`). + 1. `lora_exclude_path_regex` — skip any module whose dotted path matches, applied after target matching. VLM only; needed when a vision tower shares projection names with the LLM (Cosmos3-Edge). 1. `[model.ema]` 1. `enabled`, `rate`, `iteration_shift` — Exponential moving average of generation-pathway weights. Full fine-tunes typically enable it; LoRA recipes leave it off. 1. `[model.parallelism]` diff --git a/examples/launch_sft_videophy2_lora_edge.sh b/examples/launch_sft_videophy2_lora_edge.sh new file mode 100755 index 00000000..8997627e --- /dev/null +++ b/examples/launch_sft_videophy2_lora_edge.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# Structured-TOML launch for videophy2_lora_edge (LoRA SFT on VideoPhy-2 via +# CosmosDataLoader) targeting the Cosmos3-Edge reasoner backbone (public, +# ungated nvidia/Cosmos3-Edge, model_type cosmos3_edge — native HF metadata, +# no remote code; the classes are registered in-framework). Drives +# cosmos_framework.scripts.train against +# examples/toml/sft_config/videophy2_lora_edge.toml. +# +# [job].task = "vlm" — picks cosmos_framework/configs/base/reasoner/config.py as the base config. +# +# Reasoner weights load DIRECTLY from the nvidia/Cosmos3-Edge snapshot resolved +# via model_name: the training loader follows the repo's root safetensors index +# into its weight shards. No converter step and no required weights env var. +# +# Required env: +# VIDEOPHYSICS_ROOT dir containing videophy2_train/ and videophy2_val/ +# (each with meta.json + media/ + text/). Populate via +# `python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf`. +# +# Optional env: +# VLM_SAFETENSORS_PATH local directory of reasoner safetensors to load +# INSTEAD of the nvidia/Cosmos3-Edge snapshot. +# HF_TOKEN NOT needed for nvidia/Cosmos3-Edge (the repo is ungated). +# NPROC_PER_NODE torchrun GPUs per node; default 8. Set 4 on a GB200x4 +# node — Edge is only 2B and fits a 4-GPU allocation. +# WANDB_API_KEY the TOML sets wandb_mode="online"; export a key or set +# EXTRA_TAIL_OVERRIDES='job.wandb_mode=offline'. +# EXTRA_TAIL_OVERRIDES extra Hydra-style overrides. On nodes without a +# flash-attn wheel fall back to the portable attention +# impl: +# EXTRA_TAIL_OVERRIDES='model.config.policy.attn_implementation=sdpa' +# +# Usage (from the repo root, inside the training container): +# VIDEOPHYSICS_ROOT=/path/to/videophysics bash examples/launch_sft_videophy2_lora_edge.sh +# # on a 4-GPU node (e.g. GB200x4): +# NPROC_PER_NODE=4 VIDEOPHYSICS_ROOT=/path/to/videophysics bash examples/launch_sft_videophy2_lora_edge.sh + +TOML_FILE="examples/toml/sft_config/videophy2_lora_edge.toml" + +# The base recipe enables hf_export so eval_videophy2 can read each save as HF +# safetensors. The export is correct for a LoRA run -- HFExportCallback merges +# the adapter into the base weights -- but it gathers the full backbone onto +# rank 0 and writes it out, which is wasted on a convergence smoke run. Drop +# this line when the run's output is actually meant to be evaluated. +# +# This is the ONE knob the structured TOML cannot express (no [checkpoint] +# hf_export field in the schema); everything else lives in the TOML. +TAIL_OVERRIDES=( + checkpoint.hf_export.enabled=false + ${EXTRA_TAIL_OVERRIDES:-} +) + +# Optional: when VLM_SAFETENSORS_PATH is set, plumb it to backbone.safetensors_path +# so the framework loads reasoner weights from the local directory instead of the +# nvidia/Cosmos3-Edge snapshot (the public HF model_name still drives +# tokenizer/architecture discovery). +if [[ -n "${VLM_SAFETENSORS_PATH:-}" ]]; then + TAIL_OVERRIDES+=("model.config.policy.backbone.safetensors_path=$VLM_SAFETENSORS_PATH") +fi + +source "$(dirname "${BASH_SOURCE[0]}")/_sft_launcher_common.sh" diff --git a/examples/launch_sft_videophy2_lora_super.sh b/examples/launch_sft_videophy2_lora_super.sh new file mode 100755 index 00000000..f0750702 --- /dev/null +++ b/examples/launch_sft_videophy2_lora_super.sh @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# Structured-TOML launch for videophy2_lora_super (LoRA SFT on VideoPhy-2 via +# CosmosDataLoader, Cosmos3-Super tier / Qwen3-VL-32B). Drives +# cosmos_framework.scripts.train against +# examples/toml/sft_config/videophy2_lora_super.toml. +# +# [job].task = "vlm" — picks cosmos_framework/configs/base/reasoner/config.py as the base config. +# +# Freezing the 32B backbone and training only rank-16 adapters is what makes this +# tier comfortable on a 4-GPU allocation: optimizer state is adapter-sized. +# +# Required env: +# VIDEOPHYSICS_ROOT dir containing videophy2_train/ and videophy2_val/ +# (each with meta.json + media/ + text/). Populate via +# `python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf`. +# +# Optional env: +# HF_TOKEN for gated Qwen3-VL-32B-Instruct downloads. +# VLM_SAFETENSORS_PATH local directory of pre-converted Qwen3-VL-32B safetensors. +# When unset the framework downloads the public +# Qwen/Qwen3-VL-32B-Instruct snapshot (~64 GB on first run). +# NPROC_PER_NODE torchrun GPUs per node; default 8. Set 4 on a GB200x4 node. +# WANDB_API_KEY the TOML sets wandb_mode="online"; export a key or set +# EXTRA_TAIL_OVERRIDES='job.wandb_mode=offline'. +# +# Usage (from the repo root, inside the training container): +# VIDEOPHYSICS_ROOT=/path/to/videophysics bash examples/launch_sft_videophy2_lora_super.sh +# # on a 4-GPU node (e.g. GB200x4): +# NPROC_PER_NODE=4 VIDEOPHYSICS_ROOT=/path/to/videophysics bash examples/launch_sft_videophy2_lora_super.sh + +TOML_FILE="examples/toml/sft_config/videophy2_lora_super.toml" + +# Super-variant allocator tweak: expandable_segments so the 32B backbone fits +# without OOM. (Unlike launch_sft_vision_super.sh we do NOT clear +# LD_LIBRARY_PATH — this reasoner recipe decodes VideoPhy-2 clips with torchcodec, +# which dlopen()s the CUDA NPP + FFmpeg libs off LD_LIBRARY_PATH.) +export PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}" + +# The base recipe enables hf_export so eval_videophy2 can read each save as HF +# safetensors. The export is correct for a LoRA run -- HFExportCallback merges +# the adapter into the base weights -- but it gathers a 32B backbone (~64 GB) onto +# rank 0 and writes it out, which is wasted on a convergence smoke run. Drop +# this line when the run's output is actually meant to be evaluated. +# +# This is the ONE knob the structured TOML cannot express (no [checkpoint] +# hf_export field in the schema); everything else lives in the TOML. +TAIL_OVERRIDES=( + checkpoint.hf_export.enabled=false + ${EXTRA_TAIL_OVERRIDES:-} +) + +if [[ -n "${VLM_SAFETENSORS_PATH:-}" ]]; then + TAIL_OVERRIDES+=("model.config.policy.backbone.safetensors_path=$VLM_SAFETENSORS_PATH") +fi + +source "$(dirname "${BASH_SOURCE[0]}")/_sft_launcher_common.sh" diff --git a/examples/toml/sft_config/videophy2_lora_edge.toml b/examples/toml/sft_config/videophy2_lora_edge.toml new file mode 100644 index 00000000..424c3ea5 --- /dev/null +++ b/examples/toml/sft_config/videophy2_lora_edge.toml @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# videophy2_lora_edge — LoRA SFT on VideoPhy-2 via CosmosDataLoader, on the +# Cosmos3-Edge reasoner backbone (public, ungated nvidia/Cosmos3-Edge, +# model_type cosmos3_edge — native HF metadata, no remote code). +# Base config = cosmos_framework/configs/base/reasoner/config.py (selected by [job].task="vlm"). +# +# LoRA counterpart of videophy2_sft_edge. Same dataset, dataflow, backbone and +# freeze config (the SigLIP2 tower is frozen by regex in the experiment, since +# freeze_vision_encoder=True only supports Qwen/Intern towers); the delta is the +# adapter injection plus a LoRA-scale LR. +# +# Reasoner weights load DIRECTLY from the nvidia/Cosmos3-Edge snapshot resolved +# via model_name — no converter step, no required weights env var. +# +# Dataset prep: +# python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ +# --out_root $VIDEOPHYSICS_ROOT --split both +# +# Required env at launch: VIDEOPHYSICS_ROOT (read by the experiment Python). +# +# Example launch: +# bash examples/launch_sft_videophy2_lora_edge.sh + +[job] +task = "vlm" +experiment = "videophy2_sft_edge" # the full-FT recipe; LoRA is switched on below +project = "cosmos3" +group = "vlm_videophy2_lora" +name = "videophy2_lora_edge" +wandb_mode = "online" + +[model] +# Edge delta: the "cosmos" adapter is Qwen3-VL-only and rejects the Edge +# reasoner's (cosmos3_edge) mask. flash_attention_2 is ~14% faster than sdpa with +# matching loss curves; set "sdpa" where flash-attn is not installed. +attn_implementation = "flash_attention_2" +precision = "bfloat16" + +# LoRA. Target names match the nano/super recipes, but Edge additionally needs a +# path exclusion. Verified against the LIVE module tree, not the checkpoint: +# LLM attention (modeling_cosmos3_edge.py): self_attn.{q_proj,k_proj,v_proj,o_proj} +# SigLIP2 ViT (vision_siglip2.py): self_attn.{q_proj,k_proj,v_proj,out_proj} +# Three of the four names collide, so name matching alone would also adapt the +# vision tower — which this recipe freezes (frozen_params=["model\.visual\."]). +# Adapters landing there is NOT caught by the zero-adapter assertion: they would +# still be trainable, so the run would look healthy while training the wrong +# subnetwork. lora_exclude_path_regex keeps them in the LLM. +# +# NOTE: the snapshot's model.safetensors.index.json spells the LLM projections +# to_q/to_k/to_v/to_out. Those are PRE-REMAP checkpoint keys (the loader converts +# them on load); targeting those names matches nothing in the live model. +lora_enabled = true +lora_rank = 16 +lora_alpha = 32 +lora_target_modules = "q_proj,k_proj,v_proj,o_proj" +lora_exclude_path_regex = "^model\\.visual\\." + +# Supplies arch/config/tokenizer AND weights: the training loader follows the +# snapshot's root safetensors index into its weight shards. +[model.backbone] +model_name = "nvidia/Cosmos3-Edge" + +[model.ema] +enabled = false +rate = 0.1 +iteration_shift = 0 + +[model.parallelism] +data_parallel_shard_degree = -1 # auto from WORLD_SIZE (4- or 8-GPU) +data_parallel_replicate_degree = 1 +context_parallel_shard_degree = 1 +cfg_parallel_shard_degree = 1 + +[model.compile] +enabled = false +compile_dynamic = true + +[model.activation_checkpointing] +mode = "full" +save_ops_regex = ["fmha"] +preserve_rng_state = true +determinism_check = "default" + +[optimizer] +betas = [0.9, 0.95] +eps = 1.0e-8 +fused = true +keys_to_select = ["lora_"] +# Matched to the full-FT videophy2_sft_edge recipe (lr 1e-6, wd 0.1), with lr +# lifted 5x because LoRA adapters start from lora_B=0 and only ~0.3% of params +# carry gradient. Every other field equals videophy2_sft_edge.toml. +lr = 5.0e-6 +weight_decay = 0.1 + +[scheduler] +cycle_lengths = [300] # tracks max_iter +f_max = [1.0] +f_min = [0.1] +f_start = [0.05] +verbosity_interval = 0 +warm_up_steps = [5] + +[trainer] +distributed_parallelism = "fsdp" +grad_accum_iter = 8 +logging_iter = 1 +max_iter = 300 # longer run to see the full LoRA trajectory + +[trainer.callbacks.compile_tokenizer] +compile_after_iterations = 3 +enabled = false + +[trainer.callbacks.grad_clip] +clip_norm = 1.0 +force_finite = false + +[checkpoint] +keys_to_skip_loading = [] +load_path = "???" +save_iter = 1000 + +[dataloader_train] +max_samples_per_batch = 1 +max_sequence_length = 16000 diff --git a/examples/toml/sft_config/videophy2_lora_super.toml b/examples/toml/sft_config/videophy2_lora_super.toml new file mode 100644 index 00000000..219e1868 --- /dev/null +++ b/examples/toml/sft_config/videophy2_lora_super.toml @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# videophy2_lora_super — LoRA SFT on VideoPhy-2 via CosmosDataLoader, +# Cosmos3-Super tier (Qwen3-VL-32B-Instruct). +# Base config = cosmos_framework/configs/base/reasoner/config.py (selected by [job].task="vlm"). +# +# LoRA counterpart of videophy2_sft_super. Where that recipe full-fine-tunes the +# 32B backbone, this one freezes it entirely and trains rank-16 adapters on the +# LLM attention projections — which is what makes the 32B tier comfortable on a +# 4-GPU (GB200x4) allocation: only the adapters carry optimizer state. +# +# Weights: no converter step required. With VLM_SAFETENSORS_PATH unset the +# framework loads the public Qwen/Qwen3-VL-32B-Instruct snapshot resolved from +# model_name (~64 GB download on first run). Set VLM_SAFETENSORS_PATH to use a +# local/merged snapshot instead. +# +# Want the 8B tier instead (much smaller download, good for a first run)? There +# is no separate nano LoRA recipe because it would differ from this file in two +# values only — point [job].experiment at videophy2_sft_nano and +# [model.backbone].model_name at Qwen/Qwen3-VL-8B-Instruct. Every LoRA and +# training setting below carries over unchanged. +# +# Dataset prep: +# python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ +# --out_root $VIDEOPHYSICS_ROOT --split both +# +# Required env at launch: VIDEOPHYSICS_ROOT (read by the experiment Python). +# +# Example launch: +# bash examples/launch_sft_videophy2_lora_super.sh + +[job] +task = "vlm" +experiment = "videophy2_sft_super" # the full-FT recipe; LoRA is switched on below +project = "cosmos3" +group = "vlm_videophy2_lora" +name = "videophy2_lora_super" +wandb_mode = "online" + +[model] +attn_implementation = "cosmos" +precision = "bfloat16" + +# LoRA on the 32B backbone: 64 decoder layers x 4 projections = 256 adapters. +lora_enabled = true +lora_rank = 16 +lora_alpha = 32 +lora_target_modules = "q_proj,k_proj,v_proj,o_proj" + +[model.backbone] +model_name = "Qwen/Qwen3-VL-32B-Instruct" + +[model.ema] +enabled = false +rate = 0.1 +iteration_shift = 0 + +[model.parallelism] +data_parallel_shard_degree = -1 # FSDP full shard, auto from WORLD_SIZE +data_parallel_replicate_degree = 1 +context_parallel_shard_degree = 1 # raise to 2 (needs even GPU count) if it OOMs +cfg_parallel_shard_degree = 1 + +[model.compile] +enabled = false +compile_dynamic = true + +[model.activation_checkpointing] +mode = "full" +save_ops_regex = ["fmha"] +preserve_rng_state = true +determinism_check = "default" + +[optimizer] +betas = [0.9, 0.95] +eps = 1.0e-8 +fused = true +keys_to_select = ["lora_"] +# Matched to the full-FT videophy2_sft_super recipe (lr 1e-6, wd 0.1), with lr +# lifted 5x because LoRA adapters start from lora_B=0 and only ~0.1% of params +# carry gradient. Every other field equals videophy2_sft_super.toml. +lr = 5.0e-6 +weight_decay = 0.1 + +[scheduler] +cycle_lengths = [300] # tracks max_iter +f_max = [1.0] +f_min = [0.1] +f_start = [0.05] +verbosity_interval = 0 +warm_up_steps = [5] + +[trainer] +distributed_parallelism = "fsdp" +grad_accum_iter = 8 +logging_iter = 1 +max_iter = 300 # longer run to see the full LoRA trajectory + +[trainer.callbacks.compile_tokenizer] +compile_after_iterations = 3 +enabled = false + +[trainer.callbacks.grad_clip] +clip_norm = 1.0 +force_finite = false + +[checkpoint] +keys_to_skip_loading = [] +load_path = "???" +save_iter = 1000 + +[dataloader_train] +max_samples_per_batch = 1 +max_sequence_length = 16000