NVIDIA-NeMo / NVIDIA-NeMo/RL

xtoken cross-tokenizer distillation: teacher→student logit IPC transport is node-local (no multi-node support)

Open
#2,927 1 comment 0 reactions 1 assignee Claimed by @RayenTian View on GitHub
bug Feature
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Summary
The cross-tokenizer (xtoken) off-policy distillation path moves the teacher's
full-vocab logits to the student via **CUDA IPC**. CUDA IPC handles
(`cudaIpcOpenMemHandle`) can only be opened by a process on the **same physical
node**, so the feature currently **only works within a single node**. Any layout
where a student rank must read a teacher shard produced on a different node fails.

## Symptom
On a multi-node run, validation/training aborts with:

```
torch.AcceleratorError: CUDA error: mapping of buffer object failed
(cudaErrorMapBufferObjectFailed)
```

raised from `rebuild_cuda_tensor_from_ipc` -> `torch...._new_shared_cuda` ->
`cudaIpcOpenMemHandle`.

## Root cause
The transport uses legacy per-host CUDA IPC handles (via PyTorch
`_new_shared_cuda`). These are not valid across hosts. Relevant code:

- Producer: `ensure_teacher_ipc_buffer` / `get_full_logits_ipc`
(`nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py`)
- Consumer: `rebuild_teacher_full_logits_from_ipc` /
`assemble_teacher_logits_from_shards` / `rebuild_cuda_tensor_from_ipc`
(`nemo_rl/algorithms/x_token/loss_utils.py`, `nemo_rl/models/policy/utils.py`)

## Current mitigation (PR #2745)
`assert_xtoken_ipc_node_local` (`nemo_rl/algorithms/x_token/utils.py`) fails fast
at `setup()` unless the layout keeps IPC node-local:

- single node -> always allowed;
- `num_nodes > 1` -> requires teacher and student to share `data_parallel`
degree and a node-aligned model-parallel group (`tp*cp` equal,
`<= gpus_per_node`, and dividing `gpus_per_node`).

This prevents the mid-run crash but does **not** add multi-node support.

## Repro
Run `distillation-xtoken-off-policy-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2`
(student TP4xCP2 + teacher TP2xCP2 = 8 GPUs colocated) on a cluster with 4
GPUs/node, forcing `cluster.num_nodes=2`. The teacher shards split across the 2
nodes and the student's cross-node IPC read fails.

## Proposed solutions
1. **Fabric-type CUDA IPC** (`CU_MEM_HANDLE_TYPE_FABRIC`) — works across nodes
within one NVLink/MNNVL domain (e.g. GB200 NVL72). Least architectural change
(keeps the stash + zero-copy model) but requires VMM allocation +
cuda-python/custom export-import, and only works inside an NVLink domain.
2. **NCCL P2P transport** — general (any IB/RoCE/NVLink cluster). Reuse the
existing `collect_overlapping_teacher_shards` planner, replace the IPC open
with point-to-point send/recv. Caveat: NCCL needs teacher and student to
communicate concurrently, which conflicts with the current colocated
produce-then-consume model -> likely requires a non-colocated layout.
3. **Host-staging fallback** — copy logits via pinned host memory + network.
Simple but slow (full-vocab logits are large); debug-only.

Suggested direction: make the transport pluggable (`ipc` / `fabric_ipc` /
`nccl`) with capability detection — IPC intra-node, fabric within an NVLink
domain, NCCL otherwise.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.