[model] fix: Preserve distributed context during export

[yaoyu-33]

Summary

• preserve a caller-owned distributed context in AutoBridge.export_ckpt
• retain the existing temporary Gloo lifecycle for standalone export
• add a real world-size-one Gloo regression test

Supported trigger and impact

AutoBridge.export_ckpt is a documented one-call export API. If a supported in-process caller already owns the default process group, export unconditionally tries to create another default group and fails before checkpoint loading or Hugging Face saving with:

ValueError: trying to initialize the default process group twice!

The focused validated case is world-size-one Gloo, which the downstream CPU checkpoint load path supports. This PR does not claim new multi-rank export topology support.

Root cause and fix

export_ckpt always entered temporary_distributed_context, even though adjacent import_ckpt and load_megatron_model paths reuse caller-owned distributed state. The fix selects nullcontext() when distributed is already initialized and otherwise keeps the existing temporary Gloo context.

Regression evidence

A deterministic private reproducer AST-loaded the exact production definitions and used real PyTorch Gloo while mocking only checkpoint load/save and MCore initialization beyond the failing boundary:

uv run --no-sync python repro_export_ckpt_dist.py .

• Before: exit 1 with the duplicate-default-group ValueError before load/save.
• After: exit 0; existing caller state remained initialized and load/save were reached.
• Negative control: the same run exercised standalone export and confirmed its temporary group was cleaned up.

The committed unit test encodes the caller-owned-state contract with a real world-size-one Gloo group. Local collection of the full test module was blocked before test execution by unrelated CPU optional-package imports that require CUDA; the exact-source reproducer above was used for local runtime evidence.

Validation

Passed:

uv run --no-sync python -m py_compile src/megatron/bridge/models/conversion/auto_bridge.py tests/unit_tests/models/test_auto_bridge.py
uv run --no-sync ruff check src/megatron/bridge/models/conversion/auto_bridge.py tests/unit_tests/models/test_auto_bridge.py
uv run --no-sync ruff format --check src/megatron/bridge/models/conversion/auto_bridge.py tests/unit_tests/models/test_auto_bridge.py
git diff --check
uv run --no-sync pre-commit run --all-files

No dependency, public signature, CI workflow, checkpoint-format, or multi-rank behavior changes are included.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test 6bf4f9d

[claude]

LGTM

The export_ckpt fix correctly reuses caller-owned distributed state via nullcontext() when dist.is_initialized(), avoiding a re-init crash. This matches the identical pattern already used in import_ckpt (line 1468) and export_adapter_ckpt (line 1695), so the three conversion entry points are now consistent. The new test_export_ckpt_preserves_existing_distributed_context test exercises the caller-owned path and tears the process group down in a finally.

Suggested test cases

• test_export_ckpt_preserves_existing_distributed_context (added test - verify it passes)
• No perf tests impacted.