DeepSpeed
3dd1293f - Use nvtx's module-level API so NVTX domain ranges work on older nvtx releases (#8900)

Commit
6 days ago
Use nvtx's module-level API so NVTX domain ranges work on older nvtx releases (#8900) Since #7988,  `CUDA_Accelerator` opens NVTX domain ranges with `nvtx.get_domain(name).push_range(message=...)` , which only exists in `nvtx >= 0.2.15`. With an older nvtx installed, every CUDA training run crashes on the first ` engine.forward` , at any ZeRO stage. Rather than `pinning  nvtx>=0.2.15`  and forcing users onto the newest releases, this PR passes the domain name to nvtx's module-level  `push_range / pop_range`, which work on both old and new releases (tested 0.2.10 - 0.2.16). The recorded ranges are unchanged. Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>
Author
Parents
Loading