[ramanathan831]

Automatic HF checkpoint export wrote LoRA base and adapter parameters separately, so a plain Hugging Face loader could silently use unmerged weights. This PR exports merged dense LoRA/RSLoRA weights without changing live training parameters.

• Merge detached, fully gathered tensors in at least FP32 with TF32/autocast disabled, restore the caller's CUDA matmul precision, then cast the completed result to the output dtype. Remove only adapter keys owned by native LoraInjectedLinear modules.
• Validate shared weights before parameter deduplication or tensor collectives. Reject a base shared between plain and adapted paths, or distinct adapters, in either registration order. Preserve ordinary non-LoRA tied weights and aliases of the same LoRA module.
• Gather base/A/B tensors in the same order on every rank. Give CPU snapshots independent storage so background writes cannot observe subsequent optimizer updates. Retain wrapper normalization and collision checks.
• Preserve the pre-PR implementations of training forward() and the existing in-place/offline merge APIs, including their rounding and ambient-autocast behavior. The new numerical policy applies only to automatic export.

Validation of 0b43ed1 on H100 with locked cu128-train dependencies:

• 47 focused CPU/CUDA tests passed, including local HF reloads, LoRA/RSLoRA, BF16/FP32 exports, aliases, snapshots, legacy numerical boundaries, and CUDA precision restoration on success and failure.
• Ten FSDP cases passed on each of two ranks: dense, LoRA and tied-embedding reloads with TF32 on/off, both ambiguous-alias orders, and rank-local failure propagation. Original reload tolerances are unchanged.
• Workflow-selected L0 pytest reported 781 passed, 13 skipped. Ten skips are the separately exercised FSDP cases; three are existing architecture/NATTEN guards. Isolated attention had 1 pass and 1 existing capability skip.
• CFG parallelism: 3 passed per rank on two ranks. Context parallelism: 12 passed per rank on four ranks, retaining the workflow's existing deselection.
• Changed-file Ruff/format and Pyrefly (zero errors), both full pre-commit configurations, source-hash checks, legacy-API AST comparison, and git diff --check passed.

These local runs use H100/Python 3.13.14 rather than hosted H200/Python 3.13.15. The L0 pytest report has zero failures, but its outer local shell encountered a post-test runner-edit error; subsequent invocations of the finalized runner exit successfully. See the CI investigation and complete validation limits. Hosted CI must still validate the new head.

The focused reloads use a tiny randomly initialized Qwen2 model and local safetensors. This revision does not claim new full-training, remote-upload, or downstream execution results.

Backport requirements: corrected LoRA artifacts, independent CPU snapshots, and explicit rejection of ambiguous aliases are intentional behavior changes. Preserve unrelated valid workflows, configuration defaults, training state, and the legacy merge numerical policy. This is not a zero-impact backport.

.file_mapping.json maps the two production files, but neither their public test files nor the GitHub GPU workflow. Carry equivalent coverage from hf_export_test.py and lora_test.py, adapt destination imports, and preserve existing downstream tests. Include non-LoRA exports, valid ties, alias-order rejection, snapshot independence, BF16/autocast boundaries, TF32 on/off, precision-setting restoration, and distributed gather/rejection/failure-propagation checks. Run the distributed cases explicitly; single-process pytest skips them.

python -m pytest cosmos_framework/callbacks/hf_export_test.py cosmos_framework/utils/generator/lora_test.py -o addopts= -m L0
torchrun --nproc_per_node=2 -m pytest cosmos_framework/callbacks/hf_export_test.py -o addopts= -m L0 -k fsdp

[lfengad]

Review of 7a54b49: the non-mutating automatic LoRA export fix is justified, but I found one reproducible export edge case and one broader compatibility change that should be made explicit.

1. Tied-base rejection depends on module registration order

_gather_weights traverses deduplicated named_parameters(), and the final check only notices a LoRA base that was skipped.

With a plain Linear and a LoRA Linear sharing the same base weight, the existing test covers plain-first registration and correctly raises. Reversing registration order does not raise:

base = torch.nn.Linear(4, 3, bias=False)
adapted = LoraInjectedLinear(base, rank=2, alpha=4)
network = torch.nn.ModuleDict({"adapted": adapted, "base": base})

wrapped = SimpleNamespace(model=SimpleNamespace(model=network))
_, manifest, _ = HFExportCallback(dtype="float32")._gather_weights(wrapped)
# Currently accepted; manifest contains only "adapted.weight".
# "base.weight" is silently omitted, although the plain and adapted
# paths need different effective weights when the adapter is nonzero.

Could you check aliases before parameter deduplication and reject this ambiguous plain/LoRA shared-base case in either order? Please parameterize the existing tied-base test over both orders. This is a specific shared-weight edge case, not evidence that ordinary Qwen/VLM exports universally fail; I reproduced gathering/manifest behavior, not a full HF reload of this custom topology.

2. Existing in-place/offline merge results also change

merge_adapter_() now delegates to the new helper. That changes more than automatic checkpoint export: existing callers of the in-place/offline merge APIs get different rounding and autocast behavior.

A scalar BF16 base of 1.0, FP32 A=1.0, B=0.00392, and rank/alpha both 1 produces:

• old in-place merge: 1.0;
• new in-place merge: 1.0078125.

Accumulating the completed merge at higher precision is mathematically reasonable; I am not calling that result incorrect. However, it is not output-identical to the existing API and broadens the change beyond the missing automatic-export merge.

For minimal downstream behavioral impact, could the new non-mutating helper be used for automatic export while retaining the existing merge_adapter_() behavior in this PR, with the shared numerical-policy change reviewed separately? If you prefer to retain it here, please explicitly document the existing-API output change and add old/new boundary-case coverage rather than describing it as behavior-preserving. The training forward() itself is unchanged.

3. Downstream/backport validation

Please also note that mapped production-file backports do not necessarily carry the new public test files or GitHub GPU workflow. The downstream backport should carry equivalent regression coverage, preserve existing downstream tests, and verify non-LoRA exports as well as LoRA exports. Correcting merged LoRA weights and snapshot storage is justified; it should not be described as a zero-behavior-change backport.

Independent validation: 8 LoRA utility cases and 7 CPU gather/export-guard cases passed in an isolated harness using the actual production/test functions, with package-level logging/base-class imports isolated. The alias-order counterexample and BF16 old/new difference above were also reproduced. I did not independently rerun the full HF reload suite, the two-H100 cases, remote uploads, or full training.

[lfengad]

Following up on the review above, these are the two concrete revisions requested for this PR:

1. Make LoRA shared-base handling independent of registration order. Build the alias information before parameter deduplication, e.g. using named_parameters(remove_duplicate=False) and grouping names by parameter identity. Do not rely solely on the final missing_merges check. At minimum, the ambiguous plain/LoRA shared-base case from the review must receive the same explicit handling in both orders. Ordinary valid non-LoRA tied weights must remain supported, and distributed failure handling must be consistent across ranks.

2. Keep the existing in-place/offline merge API unchanged in this narrowly scoped fix. Retain the new non-mutating merged_weight() for automatic export, but restore the pre-PR implementation of merge_adapter_() so fixing automatic export does not incidentally change existing callers' rounding/autocast behavior. If the legacy method has a separately demonstrated correctness bug, that fix is allowed, but please identify and review its affected behavior explicitly rather than bundling it as behavior-preserving cleanup.

Together with the scheduler-state fix requested on #307, these are the three concrete changes from this review.

For the I4 mapped backport, the intended boundary is: preserve unrelated existing valid workflows, configuration defaults and training state; allow the necessary differences from confirmed correctness fixes. Correct LoRA export merging and independent CPU snapshot storage can remain. There is no requirement to keep a broken export or to disable every bug fix by default. Conversely, this should not be represented as an unconditional zero-impact change: corrected LoRA exports necessarily differ from the previously incorrect artifacts. Carry equivalent downstream regression coverage with the production changes.

[ramanathan831]

Addressed both the original review and follow-up in 28e15be.

1. Registration-order independence: _lora_merge_plan() now inspects named_parameters(remove_duplicate=False) and groups aliases by parameter identity before deduplication or tensor gathering. It rejects plain/adapted shared bases in both orders, and distinct adapters sharing one base. Ordinary tied non-LoRA weights and aliases of the same LoRA module remain supported. The two-rank FSDP regressions assert identical errors on both ranks, no entry into DTensor.full_tensor(), and unchanged local shards. Valid dense, LoRA and tied-embedding exports still reload successfully. See hf_export.py:62 and hf_export_test.py:300.

2. Existing merge compatibility: restored the exact pre-PR merge_adapter_() implementation; automatic export alone uses merged_weight(). The original training forward() and all three existing merge APIs have identical ASTs to current main. Twelve boundary cases cover the layer, network and offline merge/strip APIs with BF16/FP32 bases and autocast on/off. The scalar BF16 example now remains 1.0 through each existing API, while automatic export retains 1.0078125. See lora.py:118 and lora_test.py:114.

3. Backport boundary: updated the PR description to state the intended artifact/snapshot changes, preserve unrelated valid workflows/defaults/training state, and require equivalent downstream regressions. I checked that .file_mapping.json maps the production files but not these public tests or the GPU workflow. The notes explicitly require carrying the tests and running the two-rank cases; no downstream execution is claimed.

Validation: the new CPU regressions first reproduced ten failures, and the corrected FSDP fixture reproduced both alias-order failures on both ranks before the production fixes. On the final branch, 41 focused CPU tests and all five distributed cases per H100 rank pass. The workflow-selected H100 L0 suite passed with 754 passed and 8 skipped, the isolated attention check had 1 pass and 1 capability skip (FlashAttention-4 sparse backward is unavailable), and all 172 AutoML helper tests passed. Five L0 skips are the FSDP cases run separately; the other three are existing hardware/backend capability guards. Ruff/format, both full pre-commit configurations, changed-file Pyrefly, and diff checks also pass. These runs include actual local HF reloads and NCCL/FSDP collectives; they do not claim a new full-training or remote-upload run.

[ramanathan831]

Pushed 0b43ed1 after investigating the failed GPU unit job.

The hosted log shows the first dense FSDP export test failing before a peer's NCCL watchdog abort. That abort prevented pytest from printing the original assertion. The original five cases pass under default settings on H100, so I cannot identify the exact initial H200 assertion from the available evidence. The investigation did reproduce and fix these issues:

• The reload oracle mixed CUDA and CPU numerics. Enabling TF32 makes identical, unexported dense weights fail the old GPU-versus-CPU logits comparison. The distributed export test now computes its reference and reload on CPU while retaining real CUDA FSDP gathering. The original tolerances are unchanged; exports are exercised with TF32 both enabled and disabled.
• Automatic LoRA export still lost precision under TF32. After correcting the oracle, the TF32 LoRA export continued to fail. Disabling autocast alone did not prevent truncated-input matmul. merged_weight() now temporarily requests full FP32 CUDA matmul and restores the exact previous precision setting, including on exceptions. Six CUDA LoRA/RSLoRA cases cover highest, high, and medium; four failed before this production fix and all six now pass. The existing forward() and three in-place/offline merge APIs remain AST-identical to the PR base, preserving the review's compatibility boundary.
• Rank-local save/reload failures masked themselves as NCCL timeouts. Validation tracebacks now reach every rank before the next test. Real two-rank fault injection covers an assertion on rank 0 and an I/O error on rank 1, and verifies that the group remains usable afterward.

Validation of the pushed source on H100, using the locked cu128-train dependencies:

• Focused export/LoRA suite: 47 passed; its 10 distributed cases skip in single-process pytest and pass separately below.
• FSDP export: 10 passed per rank, two ranks.
• Workflow-selected L0 pytest: 781 passed, 13 skipped. Ten skips are the separately exercised FSDP cases; three are existing architecture/NATTEN guards.
• Isolated attention: 1 passed, 1 existing capability skip.
• CFG parallelism: 3 passed per rank, two ranks. Context parallelism: 12 passed per rank, four ranks, with the workflow's existing deselection retained.
• Changed-file Ruff/format and Pyrefly, both complete pre-commit configurations, source-hash checks, and git diff --check passed.

Validation limits: this was H100 with Python 3.13.14, versus hosted H200 with 3.13.15; torch 2.10.0+cu128 and transformers 4.57.6 match. The L0 pytest run wrote a zero-failure JUnit report, but its enclosing local shell returned 2 afterward because I edited the runner script while it was running. The finalized script passes bash -n, and subsequent attention/CFG/context-parallel invocations exit successfully. I am reporting the pytest results separately from that harness error. No new full-training, remote-upload, or downstream run is claimed. The new hosted CI run is still needed to confirm the H200 gate.