DeepSpeed
17c65954 - Avoid graph break due to unsupported frozenset (#7105)

Comment changes are shownComment changes are hidden
Commit
139 days ago
Avoid graph break due to unsupported frozenset (#7105) This PR is a continuation of the efforts to improve Deepspeed performance when using PyTorch compile. The `fetch_sub_module()` routine makes use of the `frozenset` which is problematic because: 1. `iter_params` returns an iterable over model parameters 2. `frozenset` wraps this iterable, making it unmodifiable 3. PyTorch’s compilation process cannot infer how `frozenset` interacts with tensors, leading to a graph break. If we replace the `frozenset` with a modifiable `set`, then there is no longer such graph break. Signed-off-by: Max Kovalenko <mkovalenko@habana.ai> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Author
Parents
  • deepspeed/runtime/zero
    • File
      partitioned_param_coordinator.py