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>