From 8ed01b4d529e8ea5cfe4f8f22d2d772ceeb86670 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Mon, 3 Aug 2026 20:31:58 -0700 Subject: [PATCH 1/8] feat(vlm): LoRA SFT support for VLM + VideoPhy-2 LoRA recipes Extend LoRA post-training from the VFM generation pathway to the VLM (reasoner) HF backbone. Core: - PolicyConfig gains lora_enabled/lora_rank/lora_alpha/lora_target_modules/ lora_exclude_path_regex, replacing the unused `lora: str|None` field. - VLMModel._init_vlm injects adapters on the meta-device backbone BEFORE parallelize() (post-FSDP injection sizes lora_B to the per-rank shard and crashes at forward), then initializes lora_A/lora_B after materialization and weight loading, and asserts the sharded init actually landed. - LoRA-only trainability is re-asserted after _apply_freeze_config, which would otherwise unfreeze base weights by regex and silently turn a LoRA run into a partial full fine-tune. Adapter placement is logged so mistargeting is visible rather than looking healthy. - Adapter keys are added to the checkpoint skip patterns; they never exist in a pretrained checkpoint and would trip the completeness check. - _inject_lora_inplace gains exclude_path_regex: Cosmos3-Edge shares three of four projection names between its LLM and its SigLIP2 vision tower, so name matching alone cannot keep adapters out of the frozen ViT. TOML plumbing: - model.lora_* now remaps to model.config.policy.lora_* for VLM instead of being dropped; lora_exclude_path_regex is dropped for VFM only. - _hydra_format quotes any value outside Hydra's unquoted-value charset -- regex metacharacters like ^ or \ previously threw LexerNoViableAltException. ${oc.env:...} stays unquoted so it resolves. Recipes: videophy2_z_lora.py (nano/super/edge LoRA variants of the videophy2_sft_* experiments) plus matching TOML configs and launch shells. The _z_ prefix keeps the module sorting after the recipes it deepcopies -- import_all_modules_from_package reloads alphabetically, and cloning from a later-reloaded sibling breaks dataloader pickling. Co-Authored-By: Claude Opus 5 (1M context) --- .../base/reasoner/defaults/policy_config.py | 24 ++- .../reasoner/experiment/videophy2_z_lora.py | 152 ++++++++++++++++++ .../configs/toml_config/sft_config.py | 31 +++- .../configs/toml_config/toml_config_helper.py | 35 ++-- cosmos_framework/model/generator/vlm_model.py | 137 +++++++++++++++- cosmos_framework/utils/generator/lora.py | 32 +++- examples/launch_sft_videophy2_lora_edge.sh | 55 +++++++ examples/launch_sft_videophy2_lora_nano.sh | 46 ++++++ examples/launch_sft_videophy2_lora_super.sh | 50 ++++++ .../toml/sft_config/videophy2_lora_edge.toml | 126 +++++++++++++++ .../toml/sft_config/videophy2_lora_nano.toml | 119 ++++++++++++++ .../toml/sft_config/videophy2_lora_super.toml | 109 +++++++++++++ 12 files changed, 890 insertions(+), 26 deletions(-) create mode 100644 cosmos_framework/configs/base/reasoner/experiment/videophy2_z_lora.py create mode 100755 examples/launch_sft_videophy2_lora_edge.sh create mode 100755 examples/launch_sft_videophy2_lora_nano.sh create mode 100755 examples/launch_sft_videophy2_lora_super.sh create mode 100644 examples/toml/sft_config/videophy2_lora_edge.toml create mode 100644 examples/toml/sft_config/videophy2_lora_nano.toml create mode 100644 examples/toml/sft_config/videophy2_lora_super.toml diff --git a/cosmos_framework/configs/base/reasoner/defaults/policy_config.py b/cosmos_framework/configs/base/reasoner/defaults/policy_config.py index 1cabfce6..4edcb991 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/base/reasoner/experiment/videophy2_z_lora.py b/cosmos_framework/configs/base/reasoner/experiment/videophy2_z_lora.py new file mode 100644 index 00000000..d48bcc79 --- /dev/null +++ b/cosmos_framework/configs/base/reasoner/experiment/videophy2_z_lora.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""VideoPhy-2 LoRA SFT recipes — LoRA counterparts of ``videophy2_sft_*``. + +Each recipe is a ``deepcopy`` of the matching full-fine-tune experiment plus the +LoRA delta below; the dataflow, dataset, freeze config, and parallelism come +along unchanged so the only variable between a LoRA run and its full-fine-tune +baseline is the adapter. + +The delta, and why each piece is needed: + +* ``model.config.policy.lora_enabled`` — turns on the injection in + ``VLMModel._init_vlm`` (meta device, pre-FSDP). +* ``optimizer.lr`` 1e-6 -> 1e-4 — the full fine-tunes ship 1e-6; a rank-16 + adapter starting from ``lora_B=0`` needs roughly two orders of magnitude more + to move at all in a short run. +* ``optimizer.keys_to_select=["lora_"]`` — belt-and-braces with the + ``requires_grad`` enforcement in ``VLMModel.__init__``; keeps the optimizer + state at adapter size rather than allocating for the frozen backbone. +* ``checkpoint.hf_export.enabled=False`` + a save_iter past ``max_iter`` — these + are convergence smoke runs; exporting a 32B HF snapshot per save would + dominate the wall clock and fill the disk. +* short cosine schedule matched to ``max_iter`` — the base recipes' 50-step + cycle would leave the LR mid-decay at iteration 100. + +Launch via ``examples/launch_sft_videophy2_lora_{nano,super,edge}.sh``. + +Why the ``_z_`` in the filename +------------------------------- +``make_config`` calls ``import_all_modules_from_package(..., reload=True)``, and +``pkgutil.iter_modules`` walks this package in ALPHABETICAL order. Reloading a +module rebinds its module-level functions to fresh objects. So a module that +deepcopies a recipe from a sibling reloaded LATER ends up holding the sibling's +pre-reload function objects, and ``pickle`` — which the dataloader workers use — +rejects them with "it's not the same object as +``...videophy2_sft_nano.build_videophy2_local_dataset``". + +This module must therefore sort AFTER every ``videophy2_sft_*`` module it clones. +``videophy2_sft_super`` gets away with the same pattern only because "super" +happens to sort after "nano". The assertion below turns that implicit ordering +constraint into a loud failure at config-load time rather than a confusing +pickling error minutes into a run. +""" + +from __future__ import annotations + +import copy + +from hydra.core.config_store import ConfigStore + +from cosmos_framework.configs.base.reasoner.experiment.videophy2_sft_nano import videophy2_sft_nano +from cosmos_framework.configs.base.reasoner.experiment.videophy2_sft_super import videophy2_sft_super +from cosmos_framework.configs.base.reasoner.experiment.videophy2_sft_edge import videophy2_sft_edge + +cs = ConfigStore.instance() + + +def _assert_reload_order() -> None: + """Fail loudly if this module no longer sorts after the recipes it clones. + + See the module docstring: a stale cross-module function reference surfaces as + a ``_pickle.PicklingError`` from a dataloader worker, which is a long way + from its cause. Comparing the object we captured against the one currently + bound on the source module catches it here instead. + """ + import pkgutil + import os + + here = os.path.basename(__file__).removesuffix(".py") + siblings = [m.name for m in pkgutil.iter_modules([os.path.dirname(__file__)])] + cloned = [m for m in siblings if m.startswith("videophy2_sft_")] + late = [m for m in cloned if m > here] + assert not late, ( + f"{here} clones {cloned} but sorts BEFORE {late}, which are reloaded after it by " + "import_all_modules_from_package(reload=True). The cloned recipes would carry " + "stale function objects and fail to pickle in the dataloader workers. " + f"Rename this module so it sorts after {late}." + ) + + +# Qwen3-VL LLM attention projections. Matched by EXACT child-module name, so the +# vision tower (``qkv`` / ``proj`` / ``linear_fc1`` / ``linear_fc2``) is not hit. +_QWEN3_VL_TARGETS = "q_proj,k_proj,v_proj,o_proj" + + +def _lora_variant(base, *, lora_target_modules: str, exclude_path_regex: str = ""): + """Clone a full-fine-tune recipe and switch LoRA on — nothing else. + + Every training hyperparameter — lr, max_iter, scheduler (warmup / cycle / + f_min), weight_decay, betas, validation cadence, grad_accum, dataset — is + inherited UNCHANGED from ``base``. A LoRA recipe is therefore an + apples-to-apples counterpart of its full-fine-tune sibling: the only deltas + are the LoRA adapter itself and training only those adapters + (``keys_to_select=["lora_"]``). This is what lets the LoRA and full-FT curves + be compared directly under identical settings. + + The launch TOML stays authoritative and can override any inherited value; the + shipped ``videophy2_lora_nano.toml`` keeps every training field identical to + ``videophy2_sft_nano.toml``. (The edge/super LoRA TOMLs deliberately raise lr + and max_iter for a longer sweep — that lives in the TOML, not here.) + """ + item = copy.deepcopy(base) + + item.model.config.policy.lora_enabled = True + item.model.config.policy.lora_rank = 16 + item.model.config.policy.lora_alpha = 32 + item.model.config.policy.lora_target_modules = lora_target_modules + item.model.config.policy.lora_exclude_path_regex = exclude_path_regex + + item.optimizer.keys_to_select = ["lora_"] + + # Adapters, not a full HF snapshot — don't export (esp. the 32B tier). + # This is a callback toggle; it does not touch the optimization. + item.checkpoint.hf_export.enabled = False + + item.job.wandb_mode = "online" + item.job.group = "vlm_videophy2_lora" + return item + + +videophy2_lora_nano = _lora_variant(videophy2_sft_nano, lora_target_modules=_QWEN3_VL_TARGETS) +videophy2_lora_super = _lora_variant(videophy2_sft_super, lora_target_modules=_QWEN3_VL_TARGETS) + +# Edge (``cosmos3_edge``) uses the same projection names as Qwen3-VL in its LLM, +# but its SigLIP2 vision tower reuses three of them — verified against the LIVE +# module tree (modeling_cosmos3_edge.py / vision_siglip2.py), NOT the checkpoint: +# LLM attention: self_attn.{q_proj,k_proj,v_proj,o_proj} +# SigLIP2 ViT: self_attn.{q_proj,k_proj,v_proj,out_proj} +# So name matching alone would also adapt the vision tower, which this recipe +# freezes. The exclusion regex is what keeps the adapters in the LLM; without it +# the zero-adapter assertion still passes and the run silently trains the ViT. +# +# (The snapshot's safetensors index spells the LLM projections to_q/to_k/to_v/ +# to_out — those are pre-remap checkpoint keys and match nothing in the model.) +_COSMOS3_EDGE_EXCLUDE = r"^model\.visual\." + +videophy2_lora_edge = _lora_variant( + videophy2_sft_edge, + lora_target_modules=_QWEN3_VL_TARGETS, + exclude_path_regex=_COSMOS3_EDGE_EXCLUDE, +) + + +for _item in [videophy2_lora_nano, videophy2_lora_super, videophy2_lora_edge]: + experiment_name = [name.lower() for name, value in globals().items() if value is _item][0] + if "job" not in _item: + _item["job"] = dict(name=experiment_name + "_${now:%Y-%m-%d}_${now:%H-%M-%S}") + else: + _item["job"]["name"] = experiment_name + "_${now:%Y-%m-%d}_${now:%H-%M-%S}" + + cs.store(group="experiment", package="_global_", name=experiment_name, node=_item) diff --git a/cosmos_framework/configs/toml_config/sft_config.py b/cosmos_framework/configs/toml_config/sft_config.py index 04d0efbc..a62d5f96 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. " + "videophy2_lora_nano)." ), ) 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 e76970a8..305d8551 100644 --- a/cosmos_framework/model/generator/vlm_model.py +++ b/cosmos_framework/model/generator/vlm_model.py @@ -237,6 +237,96 @@ def _apply_freeze_config(model: nn.Module, model_type: str, cfg) -> int: return n +def _assert_lora_initialized(model: nn.Module) -> None: + """Fail loudly if adapter init left garbage behind. + + ``init_lora_weights_post_materialization`` runs on FSDP2-sharded params, so + every write goes through DTensor. A silent no-op there would leave whatever + ``torch.empty_like`` allocated and produce NaN losses several minutes into + the run. Check the local shard of the first adapter pair instead: + ``lora_A`` must be finite and non-zero, ``lora_B`` must be exactly zero. + + Ranks whose shard of a tensor is empty (uneven FSDP split) are skipped. + """ + from cosmos_framework.utils.generator.lora import LoraInjectedLinear + + def _local(t: torch.Tensor) -> torch.Tensor: + return t.to_local() if hasattr(t, "to_local") else t + + for name, module in model.named_modules(): + if not isinstance(module, LoraInjectedLinear): + continue + a = _local(module.lora_A.weight.detach()) + b = _local(module.lora_B.weight.detach()) + if a.numel() == 0: + continue + if not torch.isfinite(a).all(): + raise RuntimeError(f"LoRA init failed: {name}.lora_A contains non-finite values after init.") + if not a.any(): + raise RuntimeError( + f"LoRA init failed: {name}.lora_A is all-zero after init. " + "kaiming_uniform_ did not reach the sharded tensor — the adapter would never learn." + ) + if b.any(): + raise RuntimeError(f"LoRA init failed: {name}.lora_B is not zero-initialized.") + log.info(f"LoRA init verified on {name} (lora_A std={a.float().std().item():.4g}, lora_B all-zero)") + return + + +def _enforce_lora_only_trainable(model: nn.Module) -> None: + """Freeze everything except LoRA adapters, in-place. + + ``inject_lora_pre_fsdp`` already does this at injection time, but + ``_apply_freeze_config`` runs later and can flip base params back to + trainable. This re-asserts LoRA-only and logs loudly when it had to undo + something, so a mis-specified freeze config is visible rather than silently + producing a partial full fine-tune. + """ + reverted = [n for n, p in model.named_parameters() if p.requires_grad and "lora_" not in n] + if reverted: + log.warning( + f"LoRA: freeze config left {len(reverted)} non-adapter parameter tensor(s) trainable " + f"(first up to 5: {reverted[:5]}); re-freezing them. Remove `trainable_params` from the " + "freeze config if you did not intend this." + ) + + lora_numel = 0 + frozen_numel = 0 + for name, param in model.named_parameters(): + is_lora = "lora_" in name + param.requires_grad_(is_lora) + if is_lora: + lora_numel += param.numel() + else: + frozen_numel += param.numel() + + assert lora_numel > 0, ( + "LoRA is enabled but 0 adapter parameters are trainable — check " + "model.config.policy.lora_target_modules against the backbone's module names." + ) + log.info( + f"LoRA-only training: {lora_numel:,} trainable adapter params, " + f"{frozen_numel:,} frozen base params " + f"({100 * lora_numel / max(1, lora_numel + frozen_numel):.3f}% trainable)" + ) + + # Where the adapters actually landed. A non-zero adapter count is NOT enough + # to conclude the targets were right: naming differs across model families + # (Qwen3-VL puts q_proj/k_proj/v_proj/o_proj in the LLM, cosmos3_edge puts + # those same names in the SigLIP2 vision tower and uses to_q/to_k/to_v/to_out + # for the LLM). Mistargeting produces a healthy-looking run that trains the + # wrong subnetwork, so print the placement and let the reader judge. + placement: dict[str, int] = {} + for name, _ in model.named_parameters(): + if "lora_" not in name: + continue + # Collapse layer indices so 28 layers report as one bucket. + bucket = re.sub(r"\.\d+\.", ".*.", name.rsplit(".lora_", 1)[0]) + placement[bucket] = placement.get(bucket, 0) + 1 + for bucket, count in sorted(placement.items(), key=lambda kv: -kv[1]): + log.info(f"LoRA placement: {count:4d} adapter tensors under {bucket}") + + class VLMModel(ImaginaireModel): """Config-instantiable ImaginaireModel for VLM training. @@ -263,6 +353,14 @@ def __init__(self, config: VLMModelConfig, checkpoint): f"freeze config applied (model_type={self.hf_config.model_type}): {n_trainable} trainable parameter tensors" ) + # LoRA-only is authoritative over the freeze config. ``_apply_freeze_config`` + # runs AFTER ``_init_vlm`` (where the adapters were injected and every base + # param frozen), and its ``trainable_params`` branch unfreezes by regex — + # which would silently un-freeze base weights and turn a "LoRA run" into a + # partial full fine-tune. Re-assert here. + if config.policy.lora_enabled: + _enforce_lora_only_trainable(self.model.model) + dp_group = None cp_group = None if self.parallel_dims is not None: @@ -344,6 +442,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 @@ -410,6 +525,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( @@ -425,6 +547,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) ── @@ -450,7 +573,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: @@ -462,6 +585,18 @@ 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. + if policy.lora_enabled: + from cosmos_framework.utils.generator.lora import init_lora_weights_post_materialization + + init_lora_weights_post_materialization(hf_model.model) + _assert_lora_initialized(hf_model.model) + # ── i. Gradient checkpointing ── # HF backbone supports only binary on/off via gradient_checkpointing_enable, # so VLMActivationCheckpointingConfig.mode is restricted to {"full", "none"}. diff --git a/cosmos_framework/utils/generator/lora.py b/cosmos_framework/utils/generator/lora.py index 8bffe56d..d7ed9400 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 @@ -71,20 +72,33 @@ def _inject_lora_inplace( target_modules: list[str], rank: int, alpha: int, + exclude_path_regex: str | None = None, ) -> int: """Replace each ```` ``nn.Linear`` child in-place with ``LoraInjectedLinear``. Match is by exact child name (e.g., ``q_proj_moe_gen``), not substring. 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. """ target_set = set(target_modules) + exclude = re.compile(exclude_path_regex) if exclude_path_regex else None replaced = 0 - for _parent_name, parent in list(network.named_modules()): + for parent_name, parent in list(network.named_modules()): for child_name, child in list(parent.named_children()): - if child_name in target_set and isinstance(child, nn.Linear): - setattr(parent, child_name, LoraInjectedLinear(child, rank, alpha)) - replaced += 1 + if child_name not in target_set or not isinstance(child, nn.Linear): + continue + path = f"{parent_name}.{child_name}" if parent_name else child_name + if exclude is not None and exclude.search(path): + continue + setattr(parent, child_name, LoraInjectedLinear(child, rank, alpha)) + replaced += 1 return replaced @@ -94,6 +108,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. @@ -126,10 +141,15 @@ def inject_lora_pre_fsdp( 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/examples/launch_sft_videophy2_lora_edge.sh b/examples/launch_sft_videophy2_lora_edge.sh new file mode 100755 index 00000000..0e52925b --- /dev/null +++ b/examples/launch_sft_videophy2_lora_edge.sh @@ -0,0 +1,55 @@ +#!/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" + +TAIL_OVERRIDES=( + ${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_nano.sh b/examples/launch_sft_videophy2_lora_nano.sh new file mode 100755 index 00000000..f6a804af --- /dev/null +++ b/examples/launch_sft_videophy2_lora_nano.sh @@ -0,0 +1,46 @@ +#!/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_nano (LoRA SFT on VideoPhy-2 via +# CosmosDataLoader, Qwen3-VL-8B-Instruct). Drives cosmos_framework.scripts.train +# against examples/toml/sft_config/videophy2_lora_nano.toml. +# +# [job].task = "vlm" — picks cosmos_framework/configs/base/reasoner/config.py as the base config. +# +# 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-8B-Instruct downloads. +# VLM_SAFETENSORS_PATH local directory of pre-converted Qwen3-VL safetensors +# (e.g. Cosmos3-Nano LM merged with Qwen3-VL visual via +# `cosmos_framework.scripts.convert_model_to_vlm_safetensors`). +# When set, plumbed to backbone.safetensors_path via a +# tail override. When unset, the framework falls back +# to the public Qwen/Qwen3-VL-8B-Instruct HF snapshot. +# 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_nano.sh +# # on a 4-GPU node: +# NPROC_PER_NODE=4 VIDEOPHYSICS_ROOT=/path/to/videophysics bash examples/launch_sft_videophy2_lora_nano.sh + +TOML_FILE="examples/toml/sft_config/videophy2_lora_nano.toml" + +TAIL_OVERRIDES=( + ${EXTRA_TAIL_OVERRIDES:-} +) + +# When VLM_SAFETENSORS_PATH is set, plumb it to backbone.safetensors_path so the +# framework loads weights from the local snapshot while keeping the public HF +# model_name for 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..1ee3463c --- /dev/null +++ b/examples/launch_sft_videophy2_lora_super.sh @@ -0,0 +1,50 @@ +#!/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}" + +TAIL_OVERRIDES=( + ${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..22baf724 --- /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_lora_edge" +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_nano.toml b/examples/toml/sft_config/videophy2_lora_nano.toml new file mode 100644 index 00000000..7d2288e8 --- /dev/null +++ b/examples/toml/sft_config/videophy2_lora_nano.toml @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# videophy2_lora_nano — LoRA SFT on VideoPhy-2, Cosmos3-Nano (Qwen3-VL-8B). +# Base config = cosmos_framework/configs/base/reasoner/config.py (selected by [job].task="vlm"). +# +# This recipe is deliberately IDENTICAL to videophy2_sft_nano.toml except that +# LoRA is switched on: same lr (1e-6), weight_decay (0.1), betas, scheduler +# (cycle=50, warmup=5), max_iter (50), grad_accum (8), sequence length, and the +# same 32-sample example dataset. The only training-relevant deltas are the four +# lora_* keys plus optimizer.keys_to_select=["lora_"]. This makes the run an +# apples-to-apples counterpart to the full-fine-tune baseline in +# outputs/train/logs/videophy2_sft_nano_sft.log — the ONLY variable is LoRA +# on/off, so the two curves are directly comparable. +# +# (The separate videophy2_lora_nano recipe geared for a longer, higher-LR sweep +# is not this file; this one exists to answer "does LoRA behave like full-FT +# under matched settings".) +# +# Dataset prep: +# python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ +# --out_root $VIDEOPHYSICS_ROOT --split train # and again with --split val +# +# Required env at launch: VIDEOPHYSICS_ROOT. +# +# Example launch: +# bash examples/launch_sft_videophy2_lora_nano.sh + +[job] +task = "vlm" +experiment = "videophy2_lora_nano" +project = "cosmos3" +group = "vlm_videophy2_lora" +name = "videophy2_lora_nano" +wandb_mode = "online" + +[model] +attn_implementation = "cosmos" +precision = "bfloat16" # was [model.parallelism].precision + +# LoRA — the only training-relevant delta vs videophy2_sft_nano.toml. Targets are +# matched by EXACT child-module name, so the four Qwen3-VL LLM projections are +# hit and the vision tower (qkv/proj/linear_fc1/linear_fc2) is not. +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-8B-Instruct" + +[model.ema] +enabled = false +rate = 0.1 +iteration_shift = 0 + +[model.parallelism] +# Original toml ships dp_shard=8 (an 8-GPU recipe). The full-FT baseline this +# compares against actually ran on a 4-GPU node with dp_shard=4 (see its +# launch_info.yaml), giving effective batch = 4 x 1 x grad_accum(8) = 32. Match +# that as-run baseline: shard across the 4 available GPUs. +data_parallel_shard_degree = 4 +data_parallel_replicate_degree = 1 +context_parallel_shard_degree = 1 +cfg_parallel_shard_degree = 1 + +[model.compile] +enabled = false # was [model.parallelism].use_torch_compile +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 +# Train adapters only. This is the second half of the LoRA delta; every other +# optimizer field below matches videophy2_sft_nano.toml. +keys_to_select = ["lora_"] +# 5x the full-FT baseline's 1e-6. LoRA adapters start from lora_B=0 and only ~0.2% +# of params carry gradient, so a modestly higher LR is warranted; everything else +# stays at the baseline values. +lr = 5.0e-6 +weight_decay = 0.1 + +[scheduler] +cycle_lengths = [300] # tracks max_iter (else LR floors at step 50) +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 # past max_iter: skip mid-run 8B DCP writes + +[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..519c5634 --- /dev/null +++ b/examples/toml/sft_config/videophy2_lora_super.toml @@ -0,0 +1,109 @@ +# 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: unlike the nano recipe there is 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. +# +# 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_lora_super" +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 From af63cc2e0743bd2194f1523f0a38c498c93d6267 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Tue, 4 Aug 2026 00:55:29 -0700 Subject: [PATCH 2/8] feat(vlm): merge LoRA adapters into base weights on HF export HFExportCallback wrote lora_A / lora_B as extra keys alongside an untouched base weight. No HF architecture declares those keys, so from_pretrained() dropped them with a warning and handed back the untuned base model -- an export that looked complete and silently discarded everything the run trained. Nothing downstream (eval_videophy2, the diffusers converter) could consume a LoRA run's export. _gather_weights now folds each adapter in: LoraInjectedLinear.merged_weight computes W + (alpha / r) * B @ A in float32 and the adapter keys are skipped, so a LoRA export is shaped exactly like a full fine-tune's. lora_A / lora_B are all-gathered at their base weight's iteration rather than at their own, which keeps the collective order identical on every rank. The scale comes from the module's _lora_scale property, the same one forward() uses, so the two cannot drift. Accumulating in float32 matters because the delta is orders of magnitude smaller than the base weight and bfloat16 would round much of it away. Wrapper-prefix stripping moves from substring replacement to dropping whole dot-separated segments. A wrapped module's own path ENDS with the wrapper segment (layer._checkpoint_wrapped_module) and has no trailing dot, so the substring form left module paths unstripped while parameter paths were stripped -- the two never matched and the merge silently did not fire under gradient checkpointing. hf_export stays disabled on the LoRA recipes. That is now purely a cost call (~64 GB of rank-0 host RAM to snapshot the 32B tier) rather than a correctness one, and the recipe docstring says so. Co-Authored-By: Claude Opus 5 (1M context) --- cosmos_framework/callbacks/hf_export.py | 96 ++++++++- cosmos_framework/callbacks/hf_export_test.py | 183 ++++++++++++++++++ .../reasoner/experiment/videophy2_z_lora.py | 15 +- cosmos_framework/utils/generator/lora.py | 17 ++ 4 files changed, 303 insertions(+), 8 deletions(-) create mode 100644 cosmos_framework/callbacks/hf_export_test.py diff --git a/cosmos_framework/callbacks/hf_export.py b/cosmos_framework/callbacks/hf_export.py index 88c9069e..24633488 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,11 +271,21 @@ 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] = {} @@ -240,12 +309,31 @@ 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), + ) + + # 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) diff --git a/cosmos_framework/callbacks/hf_export_test.py b/cosmos_framework/callbacks/hf_export_test.py new file mode 100644 index 00000000..89353ee8 --- /dev/null +++ b/cosmos_framework/callbacks/hf_export_test.py @@ -0,0 +1,183 @@ +# 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_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/experiment/videophy2_z_lora.py b/cosmos_framework/configs/base/reasoner/experiment/videophy2_z_lora.py index d48bcc79..ed27736f 100644 --- a/cosmos_framework/configs/base/reasoner/experiment/videophy2_z_lora.py +++ b/cosmos_framework/configs/base/reasoner/experiment/videophy2_z_lora.py @@ -18,9 +18,15 @@ * ``optimizer.keys_to_select=["lora_"]`` — belt-and-braces with the ``requires_grad`` enforcement in ``VLMModel.__init__``; keeps the optimizer state at adapter size rather than allocating for the frozen backbone. -* ``checkpoint.hf_export.enabled=False`` + a save_iter past ``max_iter`` — these - are convergence smoke runs; exporting a 32B HF snapshot per save would - dominate the wall clock and fill the disk. +* ``checkpoint.hf_export.enabled=False`` — a cost decision, not a correctness + one: ``HFExportCallback`` merges the adapter into the base weights, so a LoRA + export is a plain HF checkpoint just like a full fine-tune's. But these are + convergence smoke runs, and gathering a 32B backbone onto rank 0 to write a + full snapshot would cost ~64 GB of host RAM and disk for a result nobody + reads. Flip it back on per tier when a run's output is meant to be evaluated. + Note this is not the same as skipping the checkpoint: a save_iter past + ``max_iter`` only skips the mid-run writes, since the trainer force-saves once + at train end. * short cosine schedule matched to ``max_iter`` — the base recipes' 50-step cycle would leave the LR mid-decay at iteration 100. @@ -110,7 +116,8 @@ def _lora_variant(base, *, lora_target_modules: str, exclude_path_regex: str = " item.optimizer.keys_to_select = ["lora_"] - # Adapters, not a full HF snapshot — don't export (esp. the 32B tier). + # Skip the full HF snapshot on smoke runs (esp. the 32B tier). The export + # would be correct — HFExportCallback merges the adapter in — just expensive. # This is a callback toggle; it does not touch the optimization. item.checkpoint.hf_export.enabled = False diff --git a/cosmos_framework/utils/generator/lora.py b/cosmos_framework/utils/generator/lora.py index d7ed9400..ccae9603 100644 --- a/cosmos_framework/utils/generator/lora.py +++ b/cosmos_framework/utils/generator/lora.py @@ -66,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 _inject_lora_inplace( network: nn.Module, From 66629bc2d9149683e06d9c82ef2eb31eafeafbe1 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Tue, 4 Aug 2026 01:14:38 -0700 Subject: [PATCH 3/8] refactor(vlm): drop videophy2_z_lora, drive LoRA from the TOMLs _lora_variant made nine assignments. Eight were already duplicated verbatim in the three shipped TOMLs -- the four lora_* policy fields, the edge exclude regex, keys_to_select, job.group and job.wandb_mode. Only checkpoint.hf_export.enabled had no TOML equivalent, and that is one line of TAIL_OVERRIDES in each launch shell. So the module bought nothing and cost a lot: a _z_ filename whose only job was to sort after the recipes it deepcopied, forty lines of docstring explaining why, a dead _assert_reload_order() guarding that ordering (it was defined but never called, so it would not have fired anyway), and three recipes carrying stale cross-module function references that import_all_modules_from_package(reload=True) can break at dataloader-pickle time. The TOMLs now point [job].experiment at videophy2_sft_{nano,super,edge} and switch LoRA on with overrides, which inherits the dataflow instead of cloning it. Verified behavior-preserving: the resolved config for all three tiers is byte-identical to what the deleted module produced (LazyConfig.save_yaml, ~497 lines each, diff clean). Also corrected the nano TOML header, which claimed the file was identical to videophy2_sft_nano.toml "except that LoRA is switched on" and then listed lr 1e-6, max_iter 50 and cycle 50 -- the file actually ships 5e-6, 300 and 300. It now lists the real deltas. Co-Authored-By: Claude Opus 5 (1M context) --- .../reasoner/experiment/videophy2_z_lora.py | 159 ------------------ .../configs/toml_config/sft_config.py | 2 +- examples/launch_sft_videophy2_lora_edge.sh | 9 + examples/launch_sft_videophy2_lora_nano.sh | 9 + examples/launch_sft_videophy2_lora_super.sh | 9 + .../toml/sft_config/videophy2_lora_edge.toml | 2 +- .../toml/sft_config/videophy2_lora_nano.toml | 31 ++-- .../toml/sft_config/videophy2_lora_super.toml | 2 +- 8 files changed, 49 insertions(+), 174 deletions(-) delete mode 100644 cosmos_framework/configs/base/reasoner/experiment/videophy2_z_lora.py diff --git a/cosmos_framework/configs/base/reasoner/experiment/videophy2_z_lora.py b/cosmos_framework/configs/base/reasoner/experiment/videophy2_z_lora.py deleted file mode 100644 index ed27736f..00000000 --- a/cosmos_framework/configs/base/reasoner/experiment/videophy2_z_lora.py +++ /dev/null @@ -1,159 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 - -"""VideoPhy-2 LoRA SFT recipes — LoRA counterparts of ``videophy2_sft_*``. - -Each recipe is a ``deepcopy`` of the matching full-fine-tune experiment plus the -LoRA delta below; the dataflow, dataset, freeze config, and parallelism come -along unchanged so the only variable between a LoRA run and its full-fine-tune -baseline is the adapter. - -The delta, and why each piece is needed: - -* ``model.config.policy.lora_enabled`` — turns on the injection in - ``VLMModel._init_vlm`` (meta device, pre-FSDP). -* ``optimizer.lr`` 1e-6 -> 1e-4 — the full fine-tunes ship 1e-6; a rank-16 - adapter starting from ``lora_B=0`` needs roughly two orders of magnitude more - to move at all in a short run. -* ``optimizer.keys_to_select=["lora_"]`` — belt-and-braces with the - ``requires_grad`` enforcement in ``VLMModel.__init__``; keeps the optimizer - state at adapter size rather than allocating for the frozen backbone. -* ``checkpoint.hf_export.enabled=False`` — a cost decision, not a correctness - one: ``HFExportCallback`` merges the adapter into the base weights, so a LoRA - export is a plain HF checkpoint just like a full fine-tune's. But these are - convergence smoke runs, and gathering a 32B backbone onto rank 0 to write a - full snapshot would cost ~64 GB of host RAM and disk for a result nobody - reads. Flip it back on per tier when a run's output is meant to be evaluated. - Note this is not the same as skipping the checkpoint: a save_iter past - ``max_iter`` only skips the mid-run writes, since the trainer force-saves once - at train end. -* short cosine schedule matched to ``max_iter`` — the base recipes' 50-step - cycle would leave the LR mid-decay at iteration 100. - -Launch via ``examples/launch_sft_videophy2_lora_{nano,super,edge}.sh``. - -Why the ``_z_`` in the filename -------------------------------- -``make_config`` calls ``import_all_modules_from_package(..., reload=True)``, and -``pkgutil.iter_modules`` walks this package in ALPHABETICAL order. Reloading a -module rebinds its module-level functions to fresh objects. So a module that -deepcopies a recipe from a sibling reloaded LATER ends up holding the sibling's -pre-reload function objects, and ``pickle`` — which the dataloader workers use — -rejects them with "it's not the same object as -``...videophy2_sft_nano.build_videophy2_local_dataset``". - -This module must therefore sort AFTER every ``videophy2_sft_*`` module it clones. -``videophy2_sft_super`` gets away with the same pattern only because "super" -happens to sort after "nano". The assertion below turns that implicit ordering -constraint into a loud failure at config-load time rather than a confusing -pickling error minutes into a run. -""" - -from __future__ import annotations - -import copy - -from hydra.core.config_store import ConfigStore - -from cosmos_framework.configs.base.reasoner.experiment.videophy2_sft_nano import videophy2_sft_nano -from cosmos_framework.configs.base.reasoner.experiment.videophy2_sft_super import videophy2_sft_super -from cosmos_framework.configs.base.reasoner.experiment.videophy2_sft_edge import videophy2_sft_edge - -cs = ConfigStore.instance() - - -def _assert_reload_order() -> None: - """Fail loudly if this module no longer sorts after the recipes it clones. - - See the module docstring: a stale cross-module function reference surfaces as - a ``_pickle.PicklingError`` from a dataloader worker, which is a long way - from its cause. Comparing the object we captured against the one currently - bound on the source module catches it here instead. - """ - import pkgutil - import os - - here = os.path.basename(__file__).removesuffix(".py") - siblings = [m.name for m in pkgutil.iter_modules([os.path.dirname(__file__)])] - cloned = [m for m in siblings if m.startswith("videophy2_sft_")] - late = [m for m in cloned if m > here] - assert not late, ( - f"{here} clones {cloned} but sorts BEFORE {late}, which are reloaded after it by " - "import_all_modules_from_package(reload=True). The cloned recipes would carry " - "stale function objects and fail to pickle in the dataloader workers. " - f"Rename this module so it sorts after {late}." - ) - - -# Qwen3-VL LLM attention projections. Matched by EXACT child-module name, so the -# vision tower (``qkv`` / ``proj`` / ``linear_fc1`` / ``linear_fc2``) is not hit. -_QWEN3_VL_TARGETS = "q_proj,k_proj,v_proj,o_proj" - - -def _lora_variant(base, *, lora_target_modules: str, exclude_path_regex: str = ""): - """Clone a full-fine-tune recipe and switch LoRA on — nothing else. - - Every training hyperparameter — lr, max_iter, scheduler (warmup / cycle / - f_min), weight_decay, betas, validation cadence, grad_accum, dataset — is - inherited UNCHANGED from ``base``. A LoRA recipe is therefore an - apples-to-apples counterpart of its full-fine-tune sibling: the only deltas - are the LoRA adapter itself and training only those adapters - (``keys_to_select=["lora_"]``). This is what lets the LoRA and full-FT curves - be compared directly under identical settings. - - The launch TOML stays authoritative and can override any inherited value; the - shipped ``videophy2_lora_nano.toml`` keeps every training field identical to - ``videophy2_sft_nano.toml``. (The edge/super LoRA TOMLs deliberately raise lr - and max_iter for a longer sweep — that lives in the TOML, not here.) - """ - item = copy.deepcopy(base) - - item.model.config.policy.lora_enabled = True - item.model.config.policy.lora_rank = 16 - item.model.config.policy.lora_alpha = 32 - item.model.config.policy.lora_target_modules = lora_target_modules - item.model.config.policy.lora_exclude_path_regex = exclude_path_regex - - item.optimizer.keys_to_select = ["lora_"] - - # Skip the full HF snapshot on smoke runs (esp. the 32B tier). The export - # would be correct — HFExportCallback merges the adapter in — just expensive. - # This is a callback toggle; it does not touch the optimization. - item.checkpoint.hf_export.enabled = False - - item.job.wandb_mode = "online" - item.job.group = "vlm_videophy2_lora" - return item - - -videophy2_lora_nano = _lora_variant(videophy2_sft_nano, lora_target_modules=_QWEN3_VL_TARGETS) -videophy2_lora_super = _lora_variant(videophy2_sft_super, lora_target_modules=_QWEN3_VL_TARGETS) - -# Edge (``cosmos3_edge``) uses the same projection names as Qwen3-VL in its LLM, -# but its SigLIP2 vision tower reuses three of them — verified against the LIVE -# module tree (modeling_cosmos3_edge.py / vision_siglip2.py), NOT the checkpoint: -# LLM attention: self_attn.{q_proj,k_proj,v_proj,o_proj} -# SigLIP2 ViT: self_attn.{q_proj,k_proj,v_proj,out_proj} -# So name matching alone would also adapt the vision tower, which this recipe -# freezes. The exclusion regex is what keeps the adapters in the LLM; without it -# the zero-adapter assertion still passes and the run silently trains the ViT. -# -# (The snapshot's safetensors index spells the LLM projections to_q/to_k/to_v/ -# to_out — those are pre-remap checkpoint keys and match nothing in the model.) -_COSMOS3_EDGE_EXCLUDE = r"^model\.visual\." - -videophy2_lora_edge = _lora_variant( - videophy2_sft_edge, - lora_target_modules=_QWEN3_VL_TARGETS, - exclude_path_regex=_COSMOS3_EDGE_EXCLUDE, -) - - -for _item in [videophy2_lora_nano, videophy2_lora_super, videophy2_lora_edge]: - experiment_name = [name.lower() for name, value in globals().items() if value is _item][0] - if "job" not in _item: - _item["job"] = dict(name=experiment_name + "_${now:%Y-%m-%d}_${now:%H-%M-%S}") - else: - _item["job"]["name"] = experiment_name + "_${now:%Y-%m-%d}_${now:%H-%M-%S}" - - cs.store(group="experiment", package="_global_", name=experiment_name, node=_item) diff --git a/cosmos_framework/configs/toml_config/sft_config.py b/cosmos_framework/configs/toml_config/sft_config.py index a62d5f96..c4bc12fd 100644 --- a/cosmos_framework/configs/toml_config/sft_config.py +++ b/cosmos_framework/configs/toml_config/sft_config.py @@ -352,7 +352,7 @@ class ModelConfig(BaseModel): "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. " - "videophy2_lora_nano)." + "examples/toml/sft_config/videophy2_lora_nano.toml)." ), ) lora_rank: int = Field( diff --git a/examples/launch_sft_videophy2_lora_edge.sh b/examples/launch_sft_videophy2_lora_edge.sh index 0e52925b..8997627e 100755 --- a/examples/launch_sft_videophy2_lora_edge.sh +++ b/examples/launch_sft_videophy2_lora_edge.sh @@ -40,7 +40,16 @@ 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:-} ) diff --git a/examples/launch_sft_videophy2_lora_nano.sh b/examples/launch_sft_videophy2_lora_nano.sh index f6a804af..f37c5df1 100755 --- a/examples/launch_sft_videophy2_lora_nano.sh +++ b/examples/launch_sft_videophy2_lora_nano.sh @@ -32,7 +32,16 @@ TOML_FILE="examples/toml/sft_config/videophy2_lora_nano.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 an 8B backbone (~16 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:-} ) diff --git a/examples/launch_sft_videophy2_lora_super.sh b/examples/launch_sft_videophy2_lora_super.sh index 1ee3463c..f0750702 100755 --- a/examples/launch_sft_videophy2_lora_super.sh +++ b/examples/launch_sft_videophy2_lora_super.sh @@ -39,7 +39,16 @@ TOML_FILE="examples/toml/sft_config/videophy2_lora_super.toml" # 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:-} ) diff --git a/examples/toml/sft_config/videophy2_lora_edge.toml b/examples/toml/sft_config/videophy2_lora_edge.toml index 22baf724..424c3ea5 100644 --- a/examples/toml/sft_config/videophy2_lora_edge.toml +++ b/examples/toml/sft_config/videophy2_lora_edge.toml @@ -25,7 +25,7 @@ [job] task = "vlm" -experiment = "videophy2_lora_edge" +experiment = "videophy2_sft_edge" # the full-FT recipe; LoRA is switched on below project = "cosmos3" group = "vlm_videophy2_lora" name = "videophy2_lora_edge" diff --git a/examples/toml/sft_config/videophy2_lora_nano.toml b/examples/toml/sft_config/videophy2_lora_nano.toml index 7d2288e8..ebaa60ce 100644 --- a/examples/toml/sft_config/videophy2_lora_nano.toml +++ b/examples/toml/sft_config/videophy2_lora_nano.toml @@ -4,18 +4,25 @@ # videophy2_lora_nano — LoRA SFT on VideoPhy-2, Cosmos3-Nano (Qwen3-VL-8B). # Base config = cosmos_framework/configs/base/reasoner/config.py (selected by [job].task="vlm"). # -# This recipe is deliberately IDENTICAL to videophy2_sft_nano.toml except that -# LoRA is switched on: same lr (1e-6), weight_decay (0.1), betas, scheduler -# (cycle=50, warmup=5), max_iter (50), grad_accum (8), sequence length, and the -# same 32-sample example dataset. The only training-relevant deltas are the four -# lora_* keys plus optimizer.keys_to_select=["lora_"]. This makes the run an -# apples-to-apples counterpart to the full-fine-tune baseline in -# outputs/train/logs/videophy2_sft_nano_sft.log — the ONLY variable is LoRA -# on/off, so the two curves are directly comparable. +# [job].experiment selects videophy2_sft_nano — the FULL-FINE-TUNE recipe. There +# is no separate LoRA experiment: everything a LoRA run needs is expressible as +# TOML overrides on top of the full-FT one, so the dataflow, dataset, freeze +# config, and callbacks are inherited rather than deepcopied. # -# (The separate videophy2_lora_nano recipe geared for a longer, higher-LR sweep -# is not this file; this one exists to answer "does LoRA behave like full-FT -# under matched settings".) +# Deltas vs videophy2_sft_nano.toml: +# lora_enabled/rank/alpha/target_modules the adapter itself +# optimizer.keys_to_select=["lora_"] train adapters only +# lr 1e-6 -> 5e-6 lora_B starts at zero and only ~0.2% +# of params carry gradient +# max_iter 50 -> 300 long enough to see the trajectory +# cycle_lengths 50 -> 300 tracks max_iter, else the LR floors +# at step 50 +# save_iter 100 -> 1000 skip the mid-run 8B DCP writes +# dp_shard 8 -> 4 match the 4-GPU node this ran on +# +# Everything else — weight_decay, betas, warmup, grad_accum, sequence length, +# and the same 32-sample example dataset — is unchanged, so the LoRA and full-FT +# curves stay comparable apart from the LR and horizon noted above. # # Dataset prep: # python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ @@ -28,7 +35,7 @@ [job] task = "vlm" -experiment = "videophy2_lora_nano" +experiment = "videophy2_sft_nano" # the full-FT recipe; LoRA is switched on below project = "cosmos3" group = "vlm_videophy2_lora" name = "videophy2_lora_nano" diff --git a/examples/toml/sft_config/videophy2_lora_super.toml b/examples/toml/sft_config/videophy2_lora_super.toml index 519c5634..85531c04 100644 --- a/examples/toml/sft_config/videophy2_lora_super.toml +++ b/examples/toml/sft_config/videophy2_lora_super.toml @@ -26,7 +26,7 @@ [job] task = "vlm" -experiment = "videophy2_lora_super" +experiment = "videophy2_sft_super" # the full-FT recipe; LoRA is switched on below project = "cosmos3" group = "vlm_videophy2_lora" name = "videophy2_lora_super" From 278c3ec0af1599a39df060aa36e62e475ba97255 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Tue, 4 Aug 2026 01:59:52 -0700 Subject: [PATCH 4/8] test(vlm): assert the LoRA merge completed, and verify it under FSDP2 Two guards in _gather_weights. The merge plan keys off named_modules() paths while the export loop keys off named_parameters() paths; if a wrapper this code does not strip ever desynchronizes the two, the merge silently no-ops and the export is the untuned base model. That is the failure mode an earlier revision of this branch actually shipped, so it gets an assertion rather than a comment: * every adapter the plan found must have been folded in, else abort with the unmerged names. Tracked on all ranks, after the last collective, so it cannot strand a peer mid-all-gather. * no exported key may contain "lora_" -- the invariant itself, asserted directly on rank 0 where the manifest lives. hf_export_fsdp_test.py covers what the CPU tests structurally cannot: with lora_A / lora_B as real DTensors, the export must equal a single-process export of the same model, key for key. It also wraps one projection in checkpoint_wrapper so a genuine _checkpoint_wrapped_module segment is in the tree. Verified on 2x GPU -- merged values matched the unsharded reference exactly (max |diff| = 0.0), no adapter keys, no collective hang, and the completeness guard fires under sharding too. Wired into gpu-tests.yml next to cfgp_ar and context_parallel, which have the same fixed-world-size shape. Under plain pytest all three skip. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/gpu-tests.yml | 10 ++ cosmos_framework/callbacks/hf_export.py | 33 ++++ .../callbacks/hf_export_fsdp_test.py | 162 ++++++++++++++++++ cosmos_framework/callbacks/hf_export_test.py | 42 +++++ 4 files changed, 247 insertions(+) create mode 100644 cosmos_framework/callbacks/hf_export_fsdp_test.py 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 24633488..86ca9c55 100644 --- a/cosmos_framework/callbacks/hf_export.py +++ b/cosmos_framework/callbacks/hf_export.py @@ -292,6 +292,7 @@ def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], di 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 / @@ -331,6 +332,7 @@ def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], di 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. @@ -360,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 index 89353ee8..eee8790d 100644 --- a/cosmos_framework/callbacks/hf_export_test.py +++ b/cosmos_framework/callbacks/hf_export_test.py @@ -174,6 +174,48 @@ def test_gather_weights_leaves_a_full_finetune_untouched(): 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) From e773656971be84723a4c845c639868f9a55e8cde Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Tue, 4 Aug 2026 02:29:08 -0700 Subject: [PATCH 5/8] refactor(vlm): drop the nano LoRA recipe, keep super + edge videophy2_lora_nano.toml differed from videophy2_lora_super.toml in four values: experiment, name, model_name, and dp_shard. Every LoRA and training setting was identical, so it showcased nothing the super recipe does not -- two near-duplicate files to keep in sync for one model_name. What is left each earns its place: * super -- freezing a 32B backbone and training rank-16 adapters is what makes that tier fit a 4-GPU allocation; this is the case LoRA exists for. * edge -- the only recipe needing lora_exclude_path_regex, because its LLM and SigLIP2 tower share three of four projection names. The 8B tier was the cheapest way in, so that path is documented rather than deleted: the super TOML header now spells out the two values to change to retarget it at Qwen3-VL-8B. sft_config.py's lora_enabled example repoints at the super TOML. No dangling references remain; toml_config tests pass (25) and both surviving recipes still resolve with the expected lora_* fields. Co-Authored-By: Claude Opus 5 (1M context) --- .../configs/toml_config/sft_config.py | 2 +- examples/launch_sft_videophy2_lora_nano.sh | 55 -------- .../toml/sft_config/videophy2_lora_nano.toml | 126 ------------------ .../toml/sft_config/videophy2_lora_super.toml | 14 +- 4 files changed, 11 insertions(+), 186 deletions(-) delete mode 100755 examples/launch_sft_videophy2_lora_nano.sh delete mode 100644 examples/toml/sft_config/videophy2_lora_nano.toml diff --git a/cosmos_framework/configs/toml_config/sft_config.py b/cosmos_framework/configs/toml_config/sft_config.py index c4bc12fd..ee55f479 100644 --- a/cosmos_framework/configs/toml_config/sft_config.py +++ b/cosmos_framework/configs/toml_config/sft_config.py @@ -352,7 +352,7 @@ class ModelConfig(BaseModel): "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_nano.toml)." + "examples/toml/sft_config/videophy2_lora_super.toml)." ), ) lora_rank: int = Field( diff --git a/examples/launch_sft_videophy2_lora_nano.sh b/examples/launch_sft_videophy2_lora_nano.sh deleted file mode 100755 index f37c5df1..00000000 --- a/examples/launch_sft_videophy2_lora_nano.sh +++ /dev/null @@ -1,55 +0,0 @@ -#!/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_nano (LoRA SFT on VideoPhy-2 via -# CosmosDataLoader, Qwen3-VL-8B-Instruct). Drives cosmos_framework.scripts.train -# against examples/toml/sft_config/videophy2_lora_nano.toml. -# -# [job].task = "vlm" — picks cosmos_framework/configs/base/reasoner/config.py as the base config. -# -# 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-8B-Instruct downloads. -# VLM_SAFETENSORS_PATH local directory of pre-converted Qwen3-VL safetensors -# (e.g. Cosmos3-Nano LM merged with Qwen3-VL visual via -# `cosmos_framework.scripts.convert_model_to_vlm_safetensors`). -# When set, plumbed to backbone.safetensors_path via a -# tail override. When unset, the framework falls back -# to the public Qwen/Qwen3-VL-8B-Instruct HF snapshot. -# 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_nano.sh -# # on a 4-GPU node: -# NPROC_PER_NODE=4 VIDEOPHYSICS_ROOT=/path/to/videophysics bash examples/launch_sft_videophy2_lora_nano.sh - -TOML_FILE="examples/toml/sft_config/videophy2_lora_nano.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 an 8B backbone (~16 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:-} -) - -# When VLM_SAFETENSORS_PATH is set, plumb it to backbone.safetensors_path so the -# framework loads weights from the local snapshot while keeping the public HF -# model_name for 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/toml/sft_config/videophy2_lora_nano.toml b/examples/toml/sft_config/videophy2_lora_nano.toml deleted file mode 100644 index ebaa60ce..00000000 --- a/examples/toml/sft_config/videophy2_lora_nano.toml +++ /dev/null @@ -1,126 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 - -# videophy2_lora_nano — LoRA SFT on VideoPhy-2, Cosmos3-Nano (Qwen3-VL-8B). -# Base config = cosmos_framework/configs/base/reasoner/config.py (selected by [job].task="vlm"). -# -# [job].experiment selects videophy2_sft_nano — the FULL-FINE-TUNE recipe. There -# is no separate LoRA experiment: everything a LoRA run needs is expressible as -# TOML overrides on top of the full-FT one, so the dataflow, dataset, freeze -# config, and callbacks are inherited rather than deepcopied. -# -# Deltas vs videophy2_sft_nano.toml: -# lora_enabled/rank/alpha/target_modules the adapter itself -# optimizer.keys_to_select=["lora_"] train adapters only -# lr 1e-6 -> 5e-6 lora_B starts at zero and only ~0.2% -# of params carry gradient -# max_iter 50 -> 300 long enough to see the trajectory -# cycle_lengths 50 -> 300 tracks max_iter, else the LR floors -# at step 50 -# save_iter 100 -> 1000 skip the mid-run 8B DCP writes -# dp_shard 8 -> 4 match the 4-GPU node this ran on -# -# Everything else — weight_decay, betas, warmup, grad_accum, sequence length, -# and the same 32-sample example dataset — is unchanged, so the LoRA and full-FT -# curves stay comparable apart from the LR and horizon noted above. -# -# Dataset prep: -# python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ -# --out_root $VIDEOPHYSICS_ROOT --split train # and again with --split val -# -# Required env at launch: VIDEOPHYSICS_ROOT. -# -# Example launch: -# bash examples/launch_sft_videophy2_lora_nano.sh - -[job] -task = "vlm" -experiment = "videophy2_sft_nano" # the full-FT recipe; LoRA is switched on below -project = "cosmos3" -group = "vlm_videophy2_lora" -name = "videophy2_lora_nano" -wandb_mode = "online" - -[model] -attn_implementation = "cosmos" -precision = "bfloat16" # was [model.parallelism].precision - -# LoRA — the only training-relevant delta vs videophy2_sft_nano.toml. Targets are -# matched by EXACT child-module name, so the four Qwen3-VL LLM projections are -# hit and the vision tower (qkv/proj/linear_fc1/linear_fc2) is not. -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-8B-Instruct" - -[model.ema] -enabled = false -rate = 0.1 -iteration_shift = 0 - -[model.parallelism] -# Original toml ships dp_shard=8 (an 8-GPU recipe). The full-FT baseline this -# compares against actually ran on a 4-GPU node with dp_shard=4 (see its -# launch_info.yaml), giving effective batch = 4 x 1 x grad_accum(8) = 32. Match -# that as-run baseline: shard across the 4 available GPUs. -data_parallel_shard_degree = 4 -data_parallel_replicate_degree = 1 -context_parallel_shard_degree = 1 -cfg_parallel_shard_degree = 1 - -[model.compile] -enabled = false # was [model.parallelism].use_torch_compile -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 -# Train adapters only. This is the second half of the LoRA delta; every other -# optimizer field below matches videophy2_sft_nano.toml. -keys_to_select = ["lora_"] -# 5x the full-FT baseline's 1e-6. LoRA adapters start from lora_B=0 and only ~0.2% -# of params carry gradient, so a modestly higher LR is warranted; everything else -# stays at the baseline values. -lr = 5.0e-6 -weight_decay = 0.1 - -[scheduler] -cycle_lengths = [300] # tracks max_iter (else LR floors at step 50) -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 # past max_iter: skip mid-run 8B DCP writes - -[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 index 85531c04..219e1868 100644 --- a/examples/toml/sft_config/videophy2_lora_super.toml +++ b/examples/toml/sft_config/videophy2_lora_super.toml @@ -10,10 +10,16 @@ # LLM attention projections — which is what makes the 32B tier comfortable on a # 4-GPU (GB200x4) allocation: only the adapters carry optimizer state. # -# Weights: unlike the nano recipe there is 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. +# 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 \ From 5e65caec9fa7686443615a54b212f08388b64613 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Tue, 4 Aug 2026 02:37:21 -0700 Subject: [PATCH 6/8] docs: cover VLM LoRA, and drop the now-false "VFM only" on lora_* Two statements were true before this branch and are not anymore: * sft_config.md's lora_enabled row said "VFM only", and its remap table listed model.lora_* among the keys skipped on the VLM path. They now remap to model.config.policy.lora_*, so lora_* is split into its own rows. * training.md's TOML reference repeated the same "VFM only". lora_target_modules was also described as matching "substrings of param names". It never did -- matching is by exact child name, or by full-path suffix when the selector contains a "." -- and the difference decides whether a selector hits one tower or both, so the row now says what actually happens and notes the VLM recipes use q_proj,k_proj,v_proj,o_proj. Added: a lora_exclude_path_regex row (nothing documented it), a training.md recipe section for the two VideoPhy-2 LoRA recipes, their launch shells in the launcher table, and the 8B retarget note in place of the deleted nano recipe. The recipe section also records that a LoRA export is merged -- an ordinary HF checkpoint with no lora_* keys -- since that is the non-obvious part of turning hf_export back on. faq.md's "enable LoRA" lever no longer claims it is a generator-only knob. Verified: no stale VFM-only claims remain, every launch shell and TOML path referenced across docs/ exists, tables are non-ragged and
tags balance. Co-Authored-By: Claude Opus 5 (1M context) --- docs/faq.md | 2 +- docs/sft_config.md | 9 ++++++--- docs/training.md | 36 +++++++++++++++++++++++++++++++++++- 3 files changed, 42 insertions(+), 5 deletions(-) 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..2e268fc1 100644 --- a/docs/sft_config.md +++ b/docs/sft_config.md @@ -89,10 +89,11 @@ Top-level model knobs. Lands at `model.config.*` on VFM and on VLM; sub-tree pat | `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_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 substrings of param names that receive an adapter. Default targets the four MoE-gen projection matrices. | +| `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]` @@ -291,7 +292,9 @@ The same TOML key lands at different Hydra paths depending on `[job].task`: | `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)* | +| `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)* | diff --git a/docs/training.md b/docs/training.md index 854a6fce..c712ddd5 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 @@ -244,6 +275,8 @@ Each launcher's default paths come from the `DATASET_PATH` + `BASE_CHECKPOINT_PA | `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]` From 43f9ee5711490810ae5daa1c9752e157103d6b44 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Tue, 4 Aug 2026 02:43:55 -0700 Subject: [PATCH 7/8] docs: align table columns to satisfy rumdl-fmt The rows added for lora_exclude_path_regex and the LoRA launch shells did not match the surrounding column padding, so the rumdl-fmt pre-commit hook rewrote them and CI failed on the modified files. Ran rumdl 0.1.62 (the pinned rev) over docs/; it now reports no issues across all 14 files. Co-Authored-By: Claude Opus 5 (1M context) --- docs/sft_config.md | 56 +++++++++++++++++++++++----------------------- docs/training.md | 28 +++++++++++------------ 2 files changed, 42 insertions(+), 42 deletions(-) diff --git a/docs/sft_config.md b/docs/sft_config.md index 2e268fc1..76b81bae 100644 --- a/docs/sft_config.md +++ b/docs/sft_config.md @@ -84,17 +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 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"`. | +| 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.) | +| `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]` @@ -280,25 +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}` | passes through | *(skipped — VFM-only)* | -| `model.lora_{enabled,rank,alpha,target_modules}` | passes through | `model.config.policy.lora_*` | +| 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)* | +| `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 c712ddd5..90cb9a6a 100644 --- a/docs/training.md +++ b/docs/training.md @@ -176,9 +176,9 @@ backbone-sized, which is what lets the 32B Super tier sit comfortably on a 4-GPU 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. | +| 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 @@ -266,17 +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_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`) | +| 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. From e988fbfcb2afe8fc05677598398529ca5279981f Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Tue, 4 Aug 2026 03:02:13 -0700 Subject: [PATCH 8/8] refactor(vlm): drop the LoRA defensive guards from vlm_model Removes _assert_lora_initialized and _enforce_lora_only_trainable (99 lines), leaving vlm_model's LoRA support as the three things that actually do work: pre-FSDP injection on meta device, post-materialization init, and the checkpoint skip patterns for adapter keys. Neither guard was load-bearing for the shipped recipes: _enforce_lora_only_trainable looked like a mutation but was a no-op repeat. inject_lora_pre_fsdp already ends with exactly the same loop -- requires_grad_ True on lora_*, False on everything else -- and logs the same trainable/frozen breakdown. It could only matter if _apply_freeze_config un-froze base weights afterwards, which only its trainable_params branch does; videophy2_sft_nano uses the named flags and videophy2_sft_edge uses frozen_params, and both of those branches only ever set requires_grad False. Verified rather than argued: injecting into an Edge-shaped tree and then applying each recipe's freeze config leaves 16 trainable adapter tensors and 0 trainable non-adapter tensors in all three states -- after injection, after edge's frozen_params, and after nano's freeze_vision_encoder. _assert_lora_initialized and the adapter-placement logging were pure diagnostics with no effect on the run. Kept: the merge-completeness and adapter-leak assertions in hf_export. Those are not the same kind of thing -- they guard a silent-wrong-output path (an export that looks complete but is the untuned base model), and one of them caught a real path-stripping bug in this branch. 56 passed / 4 skipped (callbacks + toml_config). Co-Authored-By: Claude Opus 5 (1M context) --- cosmos_framework/model/generator/vlm_model.py | 99 ------------------- 1 file changed, 99 deletions(-) diff --git a/cosmos_framework/model/generator/vlm_model.py b/cosmos_framework/model/generator/vlm_model.py index edecbd9f..c0cecb99 100644 --- a/cosmos_framework/model/generator/vlm_model.py +++ b/cosmos_framework/model/generator/vlm_model.py @@ -267,96 +267,6 @@ def _apply_freeze_config(model: nn.Module, model_type: str, cfg) -> int: return n -def _assert_lora_initialized(model: nn.Module) -> None: - """Fail loudly if adapter init left garbage behind. - - ``init_lora_weights_post_materialization`` runs on FSDP2-sharded params, so - every write goes through DTensor. A silent no-op there would leave whatever - ``torch.empty_like`` allocated and produce NaN losses several minutes into - the run. Check the local shard of the first adapter pair instead: - ``lora_A`` must be finite and non-zero, ``lora_B`` must be exactly zero. - - Ranks whose shard of a tensor is empty (uneven FSDP split) are skipped. - """ - from cosmos_framework.utils.generator.lora import LoraInjectedLinear - - def _local(t: torch.Tensor) -> torch.Tensor: - return t.to_local() if hasattr(t, "to_local") else t - - for name, module in model.named_modules(): - if not isinstance(module, LoraInjectedLinear): - continue - a = _local(module.lora_A.weight.detach()) - b = _local(module.lora_B.weight.detach()) - if a.numel() == 0: - continue - if not torch.isfinite(a).all(): - raise RuntimeError(f"LoRA init failed: {name}.lora_A contains non-finite values after init.") - if not a.any(): - raise RuntimeError( - f"LoRA init failed: {name}.lora_A is all-zero after init. " - "kaiming_uniform_ did not reach the sharded tensor — the adapter would never learn." - ) - if b.any(): - raise RuntimeError(f"LoRA init failed: {name}.lora_B is not zero-initialized.") - log.info(f"LoRA init verified on {name} (lora_A std={a.float().std().item():.4g}, lora_B all-zero)") - return - - -def _enforce_lora_only_trainable(model: nn.Module) -> None: - """Freeze everything except LoRA adapters, in-place. - - ``inject_lora_pre_fsdp`` already does this at injection time, but - ``_apply_freeze_config`` runs later and can flip base params back to - trainable. This re-asserts LoRA-only and logs loudly when it had to undo - something, so a mis-specified freeze config is visible rather than silently - producing a partial full fine-tune. - """ - reverted = [n for n, p in model.named_parameters() if p.requires_grad and "lora_" not in n] - if reverted: - log.warning( - f"LoRA: freeze config left {len(reverted)} non-adapter parameter tensor(s) trainable " - f"(first up to 5: {reverted[:5]}); re-freezing them. Remove `trainable_params` from the " - "freeze config if you did not intend this." - ) - - lora_numel = 0 - frozen_numel = 0 - for name, param in model.named_parameters(): - is_lora = "lora_" in name - param.requires_grad_(is_lora) - if is_lora: - lora_numel += param.numel() - else: - frozen_numel += param.numel() - - assert lora_numel > 0, ( - "LoRA is enabled but 0 adapter parameters are trainable — check " - "model.config.policy.lora_target_modules against the backbone's module names." - ) - log.info( - f"LoRA-only training: {lora_numel:,} trainable adapter params, " - f"{frozen_numel:,} frozen base params " - f"({100 * lora_numel / max(1, lora_numel + frozen_numel):.3f}% trainable)" - ) - - # Where the adapters actually landed. A non-zero adapter count is NOT enough - # to conclude the targets were right: naming differs across model families - # (Qwen3-VL puts q_proj/k_proj/v_proj/o_proj in the LLM, cosmos3_edge puts - # those same names in the SigLIP2 vision tower and uses to_q/to_k/to_v/to_out - # for the LLM). Mistargeting produces a healthy-looking run that trains the - # wrong subnetwork, so print the placement and let the reader judge. - placement: dict[str, int] = {} - for name, _ in model.named_parameters(): - if "lora_" not in name: - continue - # Collapse layer indices so 28 layers report as one bucket. - bucket = re.sub(r"\.\d+\.", ".*.", name.rsplit(".lora_", 1)[0]) - placement[bucket] = placement.get(bucket, 0) + 1 - for bucket, count in sorted(placement.items(), key=lambda kv: -kv[1]): - log.info(f"LoRA placement: {count:4d} adapter tensors under {bucket}") - - class VLMModel(ImaginaireModel): """Config-instantiable ImaginaireModel for VLM training. @@ -395,14 +305,6 @@ def __init__(self, config: VLMModelConfig, checkpoint): f"freeze config applied (model_type={self.hf_config.model_type}): {n_trainable} trainable parameter tensors" ) - # LoRA-only is authoritative over the freeze config. ``_apply_freeze_config`` - # runs AFTER ``_init_vlm`` (where the adapters were injected and every base - # param frozen), and its ``trainable_params`` branch unfreezes by regex — - # which would silently un-freeze base weights and turn a "LoRA run" into a - # partial full fine-tune. Re-assert here. - if config.policy.lora_enabled: - _enforce_lora_only_trainable(self.model.model) - dp_group = None cp_group = None if self.parallel_dims is not None: @@ -657,7 +559,6 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: from cosmos_framework.utils.generator.lora import init_lora_weights_post_materialization init_lora_weights_post_materialization(hf_model.model) - _assert_lora_initialized(hf_model.model) # ── h. Load the immutable standalone Parakeet artifact ── # This runs for both fresh starts and DCP resumes. On resume, DCP may