feat(vlm): LoRA SFT support for VLM + VideoPhy-2 LoRA recipes - #155
Open
lfengad wants to merge 10 commits into
Open
feat(vlm): LoRA SFT support for VLM + VideoPhy-2 LoRA recipes#155lfengad wants to merge 10 commits into
lfengad wants to merge 10 commits into
Conversation
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
_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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
Two conflicts, both in files this branch also touches. lora.py — main added _target_matches(): a target containing "." matches by full-path suffix, otherwise by exact leaf name. This branch added exclude_path_regex. They looked like competing answers to the same problem (leaf names shared across towers) but they are not: one SELECTS, the other EXCLUDES, and Cosmos3-Edge needs both. Its LLM and SigLIP2 tower share three of four projection names, so no target spelling separates them -- but the tower is one contiguous subtree, which a regex removes cleanly. Resolved by composing: match with _target_matches, then drop anything the regex hits. The docstring that merged cleanly already described both, which is the shape this restores. Verified on a synthetic Edge-shaped tree (LLM + SigLIP2 sharing q/k/v, plus a tower-shared mlp_moe_gen.up_proj): * leaf targets, no exclusion -> adapters land in BOTH towers (the bug) * leaf targets + exclusion -> 4 adapters, none in model.visual.* * path-qualified target -> selects only mlp_moe_gen.up_proj * both together -> only the LLM's q_proj vlm_model.py — purely additive on both sides: this branch's "g.3 initialize LoRA adapters", main's "h load the Parakeet artifact". Kept both in their own declared order. The Parakeet load writes only into sound_und_model.encoder and its checkpoint has no lora_* keys, so it cannot clobber the adapter init that now precedes it; noted inline. Post-merge: 57 passed / 4 skipped (callbacks + toml_config), all three LoRA TOMLs still resolve with the expected lora_* and edge exclude regex, and the 2-rank FSDP2 export test still passes on GPU. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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 <details> tags
balance.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Extend LoRA post-training from the VFM generation pathway to the VLM (reasoner) HF backbone.
Core:
lora: str|Nonefield.TOML plumbing:
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.