[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, 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 on the revised branch:

• 41 focused CPU tests passed: local HF reloads, LoRA/RSLoRA, BF16/FP32 exports, tied embeddings, ambiguous aliases in both orders, snapshot independence, and numerical boundary cases through all three existing merge APIs.
• Five FSDP cases passed on each of two H100 ranks: dense, LoRA and tied-embedding reloads, plus both ambiguous-alias orders. Tests check unchanged parameters, one gather per parameter, identical ordering, and consistent rejection before tensor collectives.
• The workflow-selected H100 L0 suite passed: 754 passed, 8 skipped; the separately launched attention check had 1 pass and 1 capability skip (FlashAttention-4 sparse backward is unavailable in this torch build). Five skips are the FSDP cases exercised above; three are existing hardware/backend capability guards. All 172 AutoML helper tests also passed after merging current main.
• Ruff lint/format, both full pre-commit configurations, changed-file Pyrefly (zero errors), and git diff --check passed. AST comparison also confirms the existing forward/merge implementations are unchanged.

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, and distributed gather/rejection 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.