Repository navigation
Conversation
Poll the communicator's abort flag in the wait loop and leave without trapping when it is set. The flag is read every `kNumAbortCheckSteps` spins. Once the flag is set the communicator is being torn down and whatever the kernel produces from there is discarded. Avoiding the trap keeps the CUDA context usable, so a rank whose peer or NIC has gone can abort and rebuild the communicator instead of being restarted. Signed-off-by: Sunita Bhaskaran <bhasunit@amazon.com>
| // Leave without trapping: a trap corrupts the CUDA context, which is exactly what an | ||
| // abort is trying to keep intact for an in-process restart. Whatever this kernel | ||
| // produces from here is discarded along with the communicator. | ||
| if (abort_flag != nullptr and ++steps >= kNumAbortCheckSteps) { |
There was a problem hiding this comment.
it should be timeout based and not steps which is deduced from the GPU arch, on more powerful GPU's it will exit faster (and there is more claim why this is not good practice).
| if (abort_flag != nullptr and ++steps >= kNumAbortCheckSteps) { | ||
| steps = 0; | ||
| if (*reinterpret_cast<volatile uint32_t*>(abort_flag) != 0) | ||
| break; |
There was a problem hiding this comment.
My main concern isn't the cost — it's what happens once we let the kernel keep executing past the wait. The moment we continue with data that never arrived, every downstream step can misbehave in its own way, and each one spawns its own failure and log path. It's an open-ended rabbit hole: we'd be signing up to anticipate and handle an effectively unbounded set of error types, with no way to prove we've covered them or that continued execution is actually safe. I don't think we can bound that surface.
My secondary concern is cost. Each wait only bails after steps >= kNumAbortCheckSteps spins, and we hit these waits a lot — not once per token, but roughly (8K / channels / batch_size) * topk times, still a high count. Every one has to burn through its abort-check interval before breaking, so the unwind is slow. An unordered shutdown would very likely be much faster.
And supporting a clean, ordered shutdown isn't trivial — lots of corner cases — for something that in the worst case fires maybe once a day. The RoI doesn't justify it.
What I do think we need: detect that the timeout was caused by nccl_dev_comm.abortFlag (as opposed to a real hang) — but stop there. The kernel shouldn't try to recover the system cleanly; the training / ML framework should own the recovery.
2 participants
Poll the communicator's abort flag in the wait loop and leave without trapping when it is set. The flag is read every
kNumAbortCheckStepsspins. Once the flag is set the communicator is being torn down and whatever the kernel produces from there is discarded.Avoiding the trap keeps the CUDA context usable, so a rank whose peer or NIC has gone can abort and rebuild the communicator instead of being restarted.
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.