[1/N] Add optional NCCL m2n weight sync backend

[princepride]

Summary

This PR is the first step toward the NCCL m2n rollout weight sync proposal in #52.

It adds an opt-in non-colocated weight sync backend:

--non-colocate-weight-sync-backend {broadcast,nccl-m2n}

The existing gather-broadcast path remains the default. Colocated rollout engines continue using the existing CUDA IPC path.

Changes

• Add UpdateWeightFromNcclm2n as an optional non-colocated updater.
• Add initial NCCL m2n layout analysis for common dense/MoE/replicated tensor patterns.
• Wire backend selection through the actor update-weight setup.
• Keep explicit fallback to the existing broadcast updater for unsupported layouts or unavailable runtime support.
• Add focused unit coverage for layout decisions and fallback/native selection behavior.

Stack

This is [1/N].

Follow-up PRs will cover the native ncclm2nReshardWithWindow runtime path, vLLM worker receive integration, staging/window management, and performance comparison against the legacy gather-broadcast path.

Test Plan

python -m pytest tests/unit/backends/megatron_utils/update_weight/test_nccl_m2n_layout.py tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_nccl_m2n.py

---

Supersedes #159 (re-opened from fork branch princepride:feature/nccl-xfer-weight-sync-mvp).

[gemini-code-assist]

Code Review

This pull request introduces an experimental, opt-in weight synchronization backend using NCCL Xfer (nccl-xfer) for non-colocated rollout engines, featuring a fallback mechanism to the existing broadcast backend. It includes layout analysis to classify weight tensors, a placeholder shim for the native NCCL Xfer Python bridge, and corresponding unit tests. Feedback on the changes highlights two important issues: first, named_params_and_buffers needs to be called with convert_to_global_name=True to ensure Hugging Face parameter names are passed to the layout analysis; second, torch.uint32 and torch.uint64 should be accessed safely via getattr to avoid compatibility issues on PyTorch versions older than 2.3.

[gemini-code-assist]

By default, named_params_and_buffers returns Megatron-style parameter names (e.g., containing query_key_value or dense_h_to_4h). However, analyze_nccl_xfer_layout expects Hugging Face parameter names (e.g., q_proj.weight, o_proj.weight) to correctly identify parallel layouts. Without converting to global (Hugging Face) names, the layout analysis will fail to match any parallel suffixes and incorrectly classify all sharded parameters as replicated, leading to incorrect synchronization behavior. Pass convert_to_global_name=True to retrieve the correct Hugging Face names.

Suggested change

	for name, tensor in named_params_and_buffers(self.args, self.model):
	for name, tensor in named_params_and_buffers(self.args, self.model, convert_to_global_name=True):

[gemini-code-assist]

Directly referencing torch.uint32 and torch.uint64 will raise an AttributeError on PyTorch versions older than 2.3, as these dtypes were only introduced recently. To ensure compatibility with older PyTorch environments, these should be accessed safely using getattr(torch, 'uint32', None) and getattr(torch, 'uint64', None).

Suggested change

	_SUPPORTED_DTYPES = {
	dtype
	for dtype in (
	torch.int8,
	torch.uint8,
	getattr(torch, "float8_e4m3fn", None),
	getattr(torch, "float8_e5m2", None),
	torch.float16,
	torch.bfloat16,
	torch.int32,
	torch.uint32,
	torch.float32,
	torch.int64,
	torch.uint64,
	torch.float64,
	)
	if dtype is not None
	}
	_SUPPORTED_DTYPES = {
	dtype
	for dtype in (
	torch.int8,
	torch.uint8,
	getattr(torch, "float8_e4m3fn", None),
	getattr(torch, "float8_e5m2", None),
	torch.float16,
	torch.bfloat16,
	torch.int32,
	getattr(torch, "uint32", None),
	torch.float32,
	torch.int64,
	getattr(torch, "uint64", None),
	torch.float64,
	)
	if dtype is not None
	}

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #33532259 | 📁 Comparing 85175d3 against latest (2702eef)

  🔍 Preview build  

42 files changed · ± 42 modified

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 32 more...