transformers
cdb4a7d9 - tp loading: async scatter + broadcast via mesh_scatter(async_op=True)

Commit
91 days ago
tp loading: async scatter + broadcast via mesh_scatter(async_op=True) distribute_tensor doesn't expose async — but mesh_scatter underneath does. Build a small _distribute_async() that uses it directly for sharded params and dist.broadcast(async_op=True) for replicated. Main loop keeps up to MAX_INFLIGHT=32 comms in flight, draining oldest as new ones queue. Scatter/broadcast now return Work handles; while N is in-flight on the NCCL stream, the CPU kicks off convert+scatter for N+1 and the thread pool reads N+2 from disk. No pipeline barriers. Numbers (8×B200, tp_plan=auto): Qwen2.5-14B: 25.30s → 23.87s load
Author
Parents
Loading