xtoken cross-tokenizer distillation: teacher→student logit IPC transport is node-local (no multi-node support)
- 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
Assessment
This issue has not been assessed yet.