Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

share inflight registry between PartitionedParameterCoordinators #3462

Merged
merged 7 commits into from
May 15, 2023
Prev Previous commit
Next Next commit
bound registry to model
  • Loading branch information
HeyangQin committed May 5, 2023
commit 30c91c04805292ce3702915e46eb4caa95548931
5 changes: 4 additions & 1 deletion deepspeed/runtime/zero/parameter_offload.py
Original file line number Diff line number Diff line change
Expand Up @@ -243,7 +243,10 @@ def __init__(self,
self._max_reuse_distance_in_numel = int(max_reuse_distance)
self._max_available_parameters_in_numel = int(max_live_parameters)
self.__allgather_stream = get_accelerator().Stream() if overlap_comm else get_accelerator().default_stream()
self.__inflight_param_registry = PartitionedParameterCoordinator.__InflightParamRegistry()

if not hasattr(module, "ds_inflight_param_registry"):
module.ds_inflight_param_registry = PartitionedParameterCoordinator.__InflightParamRegistry()
self.__inflight_param_registry = module.ds_inflight_param_registry

self.forward_hooks = []
self.backward_hooks = []
Expand Down