refactor: prepare deepspeed for ejection - #292
Conversation
jandom
added
the
safe-to-test
| def deepspeed_is_initialized() -> bool: | ||
| """Whether DeepSpeed distributed comm is initialized. Never imports deepspeed.""" | ||
| ds = sys.modules.get("deepspeed") | ||
| return ds is not None and ds.comm.comm.is_initialized() |
There was a problem hiding this comment.
This checks if the module is available without importing
There was a problem hiding this comment.
Do we require that ds comm is initialized for every use of deepspeed, also in simple inference workloads?
There was a problem hiding this comment.
That is a fair question...
| # Default inference kernels: cuEquivariance on CUDA, Triton on | ||
| # ROCm. DeepSpeed remains available as an explicit opt-in. | ||
| "use_deepspeed_evo_attention": False, | ||
| "use_cueq_triangle_kernels": not _is_rocm, |
There was a problem hiding this comment.
This is the most consequential change for inference. I don't think we want to rush this PR out, so keeping it as draft
| requires-python = ">=3.10" | ||
| # Upper-bounded at <3.15: TorchScript (torch.jit.script/script_method, still used by | ||
| # deepspeed) is unsupported and breaks on Python 3.15+, so 3.14 is the max runnable. | ||
| requires-python = ">=3.10,<3.15" |
There was a problem hiding this comment.
drive-by fix – we're getting all these torch.jit.script deprecation warnings, things just won't work with py315
There was a problem hiding this comment.
I think deepspeed newer version is almost free of script uses (see this commit)
There was a problem hiding this comment.
hopefully won't be an issue either way soon :P
jandom
changed the title
jandom
added
safe-to-test
| # Prevent deepspeed from doing triton matmul autotuning, which can hang if | ||
| # libaio is not installed and cause restart errors on preemption. Set here | ||
| # (rather than at import time) so it only runs when the DS kernel is used. | ||
| deepspeed.HAS_TRITON = False |
There was a problem hiding this comment.
Is there a way we could simply not do this? Or at least have a global config to allow the user to disable it.
jandom
added
safe-to-test
2 participants
Summary
Slowly moving away from deepspeed, keeping it as dep for now but making a couple of changes
Changes
Related Issues
Testing
Other Notes