Repository: microsoft/DeepSpeed
File: deepspeed/runtime/zero/parameter_offload.py (Python)

Code before:
            self.offload_device = offload_param_config.device
            self.offload_param_pin_memory = offload_param_config.pin_memory

        self._convert_to_zero_parameters(ds_config, module, mpu)

        for m in module.modules():
            _init_external_params(m)

        _inject_parameters(module, ZeROOrderedDict)

        self.param_numel_persistence_threshold = int(param_persistence_threshold)
        self.model_persistence_threshold = int(model_persistence_threshold)
        self.persistent_parameters = self.mark_persistent_parameters(self.param_numel_persistence_threshold,
                                                                     self.model_persistence_threshold)

        self._prefetch_bucket_sz = int(prefetch_bucket_size)
        self._max_reuse_distance_in_numel = int(max_reuse_distance)
        self._max_available_parameters_in_numel = int(max_live_parameters)
        self.__allgather_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream(
        ) if overlap_comm else get_accelerator().default_stream()

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

        self.fast_fetch_for_leaf_module = False

        if zero_module_granularity_threshold > 0:
            self.min_granularity_value = sys.maxsize
            self.min_granularity_layer = None
            self.granularity_info = set()
            self.z3_leaf_layers = []
            self._set_z3_leaf_modules_by_threshold(module, zero_module_granularity_threshold)
            self.fast_fetch_for_leaf_module = True

        self.param_coordinator = PartitionedParameterCoordinator(
            prefetch_bucket_sz=self._prefetch_bucket_sz,
            max_reuse_distance_in_numel=self._max_reuse_distance_in_numel,
            max_available_parameters_in_numel=self._max_available_parameters_in_numel,
            allgather_stream=self.__allgather_stream,
            inflight_param_registry=self.__inflight_param_registry,
            prefetch_nvme=self.offload_device == OffloadDeviceEnum.nvme,
            timers=self.timers,
            zero_quantized_weights=self.zero_quantized_weights,
            zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights,
            fast_fetch_for_leaf_module=self.fast_fetch_for_leaf_module)

        self.forward_hooks = []
        self.backward_hooks = []

        self.setup_zero_stage3_hooks()

Review comment (suggestion):
@inkcherry, since this great feature accelerates both fetch and release, I think it would be helpful for the naming to reflect this. The current naming suggests that only fetch is accelerated. So, can we consider more general naming, such as below?

```suggestion
        self.fast_sharding_for_leaf_module = False
```

Code after:
            self.offload_device = offload_param_config.device
            self.offload_param_pin_memory = offload_param_config.pin_memory

        self._convert_to_zero_parameters(ds_config, module, mpu)

        for m in module.modules():
            _init_external_params(m)

        _inject_parameters(module, ZeROOrderedDict)

        self.param_numel_persistence_threshold = int(param_persistence_threshold)
        self.model_persistence_threshold = int(model_persistence_threshold)
        self.persistent_parameters = self.mark_persistent_parameters(self.param_numel_persistence_threshold,
                                                                     self.model_persistence_threshold)

        self._prefetch_bucket_sz = int(prefetch_bucket_size)
        self._max_reuse_distance_in_numel = int(max_reuse_distance)
        self._max_available_parameters_in_numel = int(max_live_parameters)
        self.__allgather_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream(
        ) if overlap_comm else get_accelerator().default_stream()

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

        self.fast_sharding_for_leaf_module = False

        if zero_module_granularity_threshold > 0:
            self.min_granularity_value = sys.maxsize
            self.min_granularity_layer = None
            self.granularity_info = set()
            self.z3_leaf_layers = []
            self._set_z3_leaf_modules_by_threshold(module, zero_module_granularity_threshold)
            self.fast_sharding_for_leaf_module = True

        self.param_coordinator = PartitionedParameterCoordinator(
            prefetch_bucket_sz=self._prefetch_bucket_sz,
            max_reuse_distance_in_numel=self._max_reuse_distance_in_numel,
            max_available_parameters_in_numel=self._max_available_parameters_in_numel,
            allgather_stream=self.__allgather_stream,
            inflight_param_registry=self.__inflight_param_registry,
            prefetch_nvme=self.offload_device == OffloadDeviceEnum.nvme,
            timers=self.timers,
            zero_quantized_weights=self.zero_quantized_weights,
            zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights,
            fast_sharding_for_leaf_module=self.fast_sharding_for_leaf_module)

        self.forward_hooks = []
        self.backward_hooks = []

        self.setup_zero_stage3_hooks()