Skip to content

fix(weight-sync): load transfer backends lazily - #2597

Closed
nanjiangwill wants to merge 1 commit into
radixark:mainfrom
modal-projects:fix/lazy-weight-sync-imports
Closed

fix(weight-sync): load transfer backends lazily#2597
nanjiangwill wants to merge 1 commit into
radixark:mainfrom
modal-projects:fix/lazy-weight-sync-imports

Conversation

@nanjiangwill

@nanjiangwill nanjiangwill commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Summary

  • load each mutually exclusive weight-transfer implementation only in its existing selection branch
  • keep colocated, broadcast, disk-delta, and P2P backend imports structurally consistent
  • prevent non-P2P runs from loading Mooncake RDMA libraries

Root cause

The Megatron actor imported several transfer implementations at module load time. In particular, the P2P updater transitively imported Mooncake and loaded native RDMA libraries even when another transfer mode was configured.

Validation

  • python -m py_compile miles/backends/megatron_utils/actor.py
  • ruff check miles/backends/megatron_utils/actor.py
  • git diff --check

@nanjiangwill
nanjiangwill force-pushed the fix/lazy-weight-sync-imports branch from 4e18678 to 46aca9c Compare August 18, 2026 07:07
@nanjiangwill nanjiangwill changed the title fix(weight-sync): load optional P2P dependencies lazily fix(weight-sync): load transfer backends lazily Aug 18, 2026
@nanjiangwill
nanjiangwill marked this pull request as ready for review August 18, 2026 07:07

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@nanjiangwill
nanjiangwill force-pushed the fix/lazy-weight-sync-imports branch from b3cf82c to 57f0125 Compare August 24, 2026 05:05
@nanjiangwill

Copy link
Copy Markdown
Contributor Author

Closing as superseded by the weight-update refactor merged in #2752 and #2754. Current main selects the transfer implementation through get_weight_transfer_protocol(), where each backend is imported only inside its selected branch. The P2P path also imports Mooncake’s TransferEngine only when creating that engine. Non-P2P runs therefore no longer load the P2P/Mooncake dependency chain, and rebasing this actor-level patch would restore an ownership boundary that no longer exists.

@nanjiangwill
nanjiangwill deleted the fix/lazy-weight-sync-imports branch September 3, 2026 20:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant