run2/deepspeed/runtime/zero/stage_1_and_2.py

1# Copyright (c) Microsoft Corporation.
2# SPDX-License-Identifier: Apache-2.0
3
4# DeepSpeed Team
5
6from opguard.runtime_tracer import flush_trace_buffer
7from opguard.wrapped_ops import make_wrapper
8import operator
9opguard_wrapped__align_dense_tensors_flatten = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
10opguard_wrapped__all_norms_append = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
11opguard_wrapped__allreduced_unflatten = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
12opguard_wrapped__bit16_groups_flat_append = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
13opguard_wrapped__buf_0_empty = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
14opguard_wrapped__buf_1_empty = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
15opguard_wrapped__buf_copy_ = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
16opguard_wrapped__cpu_data_cpu = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
17opguard_wrapped__data_copy_ = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
18opguard_wrapped__data_mul_ = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
19opguard_wrapped__data_to = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
20opguard_wrapped__data_view_as = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
21opguard_wrapped__err_tensor = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
22opguard_wrapped__float_max = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
23opguard_wrapped__float_x_float = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
24opguard_wrapped__grad_slice_narrow = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
25opguard_wrapped__grads_in_partition_empty = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
26opguard_wrapped__inf_isinf = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
27opguard_wrapped__inf_or_nan_logical_or = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
28opguard_wrapped__invalid_grad_count__has_inf_or_nan = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
29opguard_wrapped__invalid_grad_count_bool = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
30opguard_wrapped__invalid_grad_count_zeros = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
31opguard_wrapped__meta_tensors_append = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
32opguard_wrapped__nan_isnan = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
33opguard_wrapped__narrow_average_tensor = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
34opguard_wrapped__new_grad_tensor_copy_ = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
35opguard_wrapped__new_grad_tensor_narrow = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
36opguard_wrapped__norm_is_inf_isinf = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
37opguard_wrapped__norm_is_nan_isnan = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
38opguard_wrapped__overflow_gpu_to = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
39opguard_wrapped__overflow_item = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
40opguard_wrapped__param_wrapper = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
41opguard_wrapped__partitions_append = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
42opguard_wrapped__single_grad_partition_to = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
43opguard_wrapped__single_grad_partition_zeros = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
44opguard_wrapped__stack_norm = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
45opguard_wrapped__tensor_div_ = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
46opguard_wrapped__tensor_flatten = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
47opguard_wrapped__total_norm_float = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
48opguard_wrapped__total_norm_logical_not = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
49opguard_wrapped__total_norm_pow = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
50opguard_wrapped__updated_params_unflatten = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
51opguard_wrapped__weights_partition_detach = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
52opguard_wrapped_scaled_global_grad_norm_truediv_prev_scale = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
53opguard_wrapped_total_norm_truediv_clip_grad = make_wrapper(tag_prefix = '/opt/opguard/OpGuard-runs/demo-case5-paired-20261019/workloads/run2/deps/deepspeed/runtime/zero/stage_1_and_2.py')
54import torch
55from deepspeed import comm as dist
56from packaging import version as pkg_version
57from collections import OrderedDict
58from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
59
60from deepspeed.runtime.base_optimizer import ZeROOptimizer
61from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler
62from deepspeed.runtime.utils import (empty_cache, see_memory_usage, inf, is_model_parallel_parameter,
63                                     align_dense_tensors, all_gather_dp_groups)
64from deepspeed.runtime.zero.config import ZeroStageEnum
65from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum
66from deepspeed.ops.adam import DeepSpeedCPUAdam
67from deepspeed.utils import logger
68from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank
69from deepspeed.moe.utils import is_moe_param
70from deepspeed.git_version_info import version
71
72from deepspeed.runtime.constants import PIPE_REPLICATED
73from deepspeed.accelerator import get_accelerator
74
75from deepspeed.checkpoint.constants import (DS_VERSION, GROUP_PADDINGS, PARTITION_COUNT, LOSS_SCALER,
76                                            SINGLE_PARTITION_OF_FP32_GROUPS, BASE_OPTIMIZER_STATE,
77                                            BASE_OPTIMIZER_STATE_STEP, CLIP_GRAD, ZERO_STAGE, PARAM_SLICE_MAPPINGS)
78from deepspeed.utils import link_hp_params, lazy_init_hp_params_optimizer_state
79from deepspeed.checkpoint import enable_universal_checkpoint
80
81from deepspeed.utils import groups
82# Toggle this to true to enable correctness test
83# with gradient partitioning and without
84pg_correctness_test = False
85
86OPTIMIZER_ALLGATHER_TIMER = 'optimizer_allgather'
87OPTIMIZER_GRADIENTS_TIMER = 'optimizer_gradients'
88OPTIMIZER_STEP_TIMER = 'optimizer_step'
89OPTIMIZER_TIMERS = [OPTIMIZER_ALLGATHER_TIMER, OPTIMIZER_GRADIENTS_TIMER, OPTIMIZER_STEP_TIMER]
90
91
92def input(msg):
93    return
94
95
96def split_half_float_double(tensors):
97    device_type = get_accelerator().device_name()
98    dtypes = [
99        "torch.{}.HalfTensor".format(device_type), "torch.{}.FloatTensor".format(device_type),
100        "torch.{}.DoubleTensor".format(device_type), "torch.{}.BFloat16Tensor".format(device_type)
101    ]
102    buckets = []
103    for i, dtype in enumerate(dtypes):
104        bucket = [t for t in tensors if t.type() == dtype]
105        if bucket:
106            buckets.append(bucket)
107    return buckets
108
109
110def isclose(a, b, rtol=1e-09, atol=0.0):
111    return abs(a - b) <= max(rtol * max(abs(a), abs(b)), atol)
112
113
114def lcm(x, y):
115    from fractions import gcd  # or can import gcd from `math` in Python 3
116    return x * y // gcd(x, y)
117
118
119def get_alignment_padding(tensor_list, alignment):
120    num_elements = sum([tensor.numel() for tensor in tensor_list])
121    remainder = num_elements % alignment
122    return (alignment - remainder) if remainder else remainder
123
124
125def print_rank_msg(msg):
126    print(f"rank {dist.get_rank()} - {msg}")
127
128
129def _get_padded_tensor(src_tensor, size):
130    if src_tensor.numel() >= size:
131        return src_tensor
132    padded_tensor = torch.zeros(size, dtype=src_tensor.dtype, device=src_tensor.device)
133    slice_tensor = torch.narrow(padded_tensor, 0, 0, src_tensor.numel())
134    slice_tensor.data.copy_(src_tensor.data)
135    return padded_tensor
136
137
138def _pad_tensor_by_size(src_tensor, pad_size, dtype, device):
139    padded_tensor = torch.zeros(src_tensor.numel() + pad_size, dtype=dtype, device=device)
140    padded_tensor.data[:src_tensor.numel()].copy_(src_tensor.data)
141    return padded_tensor
142
143
144class DeepSpeedZeroOptimizer(ZeROOptimizer):
145    """
146    DeepSpeedZeroOptimizer designed to reduce the memory footprint
147    required for training large deep learning models.
148
149    For more details please see ZeRO: Memory Optimization Towards Training A Trillion Parameter Models
150    https://arxiv.org/abs/1910.02054
151
152    For usage examples, refer to TODO: DeepSpeed Tutorial
153
154    """
155
156    def __init__(self,
157                 init_optimizer,
158                 param_names,
159                 timers,
160                 static_loss_scale=1.0,
161                 dynamic_loss_scale=False,
162                 dynamic_loss_args=None,
163                 verbose=True,
164                 contiguous_gradients=True,
165                 reduce_bucket_size=500000000,
166                 use_multi_rank_bucket_allreduce=True,
167                 allgather_bucket_size=5000000000,
168                 dp_process_group=None,
169                 expert_parallel_group=None,
170                 expert_data_parallel_group=None,
171                 reduce_scatter=True,
172                 overlap_comm=False,
173                 offload_optimizer_config=None,
174                 mpu=None,
175                 clip_grad=0.0,
176                 gradient_accumulation_dtype=torch.float32,
177                 communication_data_type=torch.float16,
178                 postscale_gradients=True,
179                 gradient_predivide_factor=1.0,
180                 gradient_accumulation_steps=1,
181                 ignore_unused_parameters=True,
182                 partition_grads=True,
183                 round_robin_gradients=False,
184                 has_moe_layers=False,
185                 fp16_master_weights_and_gradients=False,
186                 elastic_checkpoint=False):
187
188        if offload_optimizer_config is not None and offload_optimizer_config.device != OffloadDeviceEnum.none:
189            self.cpu_offload = True
190            self.cpu_offload_pin_memory = offload_optimizer_config.pin_memory
191        else:
192            self.cpu_offload = False
193            self.cpu_offload_pin_memory = False
194
195        if dist.get_rank() == 0:
196            logger.info(f"Reduce bucket size {reduce_bucket_size}")
197            logger.info(f"Allgather bucket size {allgather_bucket_size}")
198            logger.info(f"CPU Offload: {self.cpu_offload}")
199            logger.info(f'Round robin gradient partitioning: {round_robin_gradients}')
200        # The fused optimizer does all the work. We need this layer for two reason:
201        # 1. maintain same user API from apex.fp16_utils
202        # 2. keep common stuff here in case we need to add ne552w fused optimizer later
203
204        self.elastic_checkpoint = elastic_checkpoint
205        self.param_names = param_names
206        self.mpu = mpu
207        # differences from apex.fp16_utils:
208        # - assume all model params in fp16
209        # - assume all params requires grad
210        # - flat by groups, not keeping state. TODO: remove state explicitly?
211        # - master grad and unflat master weight never exist. TODO: a way to save out unflat master?
212        if not get_accelerator().is_available():
213            raise SystemError("Accelerator is not detected, cannot perform low precision training (e.g., fp16, bf16).")
214        self.optimizer = init_optimizer
215
216        # Use torch (un)flatten ops
217        self.flatten = _flatten_dense_tensors
218        self.unflatten = _unflatten_dense_tensors
219
220        # ZeRO stage 1 (False) or 2 (True)
221        self.partition_gradients = partition_grads
222        self.zero_stage_string = "ZeRO-2" if partition_grads else "ZeRO-1"
223
224        self.timers = timers
225
226        self.reduce_scatter = reduce_scatter
227
228        self.overlap_comm = overlap_comm
229
230        self.deepspeed_adam_offload = self.cpu_offload
231
232        self.device = get_accelerator().current_device_name() if not self.cpu_offload else 'cpu'
233
234        self.dp_process_group = dp_process_group
235        self.sequence_parallel_size = groups._get_sequence_parallel_world_size()
236        #expert parallel group
237        self.ep_process_group = expert_parallel_group
238
239        #data parallel group for experts
240        self.expert_dp_process_group = expert_data_parallel_group
241
242        #data parallel size for non-experts
243        dp_size = dist.get_world_size(group=self.dp_process_group)
244
245        #For MoE models this maybe different for different param group
246        #It will be modified during MoE setup later in the init
247        self.real_dp_process_group = [dp_process_group for i in range(len(self.optimizer.param_groups))]
248        self.partition_count = [dp_size for i in range(len(self.optimizer.param_groups))]
249
250        self.is_gradient_accumulation_boundary = True
251
252        # CPU-Offload requires contiguous gradients
253        self.contiguous_gradients = contiguous_gradients or self.cpu_offload
254
255        self.has_moe_layers = has_moe_layers
256        if self.has_moe_layers:
257            self._configure_moe_settings()
258        self._global_grad_norm = 0.
259
260        if mpu is None:
261            self.model_parallel_group = None
262            self.model_parallel_world_size = 1
263            self.model_parallel_rank = 0
264        else:
265            self.model_parallel_group = mpu.get_model_parallel_group()
266            self.model_parallel_world_size = mpu.get_model_parallel_world_size()
267            self.model_parallel_rank = bwc_tensor_model_parallel_rank(mpu)
268
269        self.overflow = False
270        self.clip_grad = clip_grad
271        self.communication_data_type = communication_data_type
272        self.gradient_predivide_factor = gradient_predivide_factor
273        self.postscale_gradients = postscale_gradients
274        self.gradient_accumulation_steps = gradient_accumulation_steps
275        self.micro_step_id = 0
276        self.ignore_unused_parameters = ignore_unused_parameters
277        self.round_robin_gradients = round_robin_gradients
278
279        self.extra_large_param_to_reduce = None
280        self.fp16_master_weights_and_gradients = fp16_master_weights_and_gradients
281
282        if self.fp16_master_weights_and_gradients:
283            assert self.cpu_offload and type(self.optimizer) in [DeepSpeedCPUAdam], \
284            f"fp16_master_and_gradients requires optimizer to support keeping fp16 master and gradients while keeping the optimizer states in fp32."\
285            f"Currently only supported using ZeRO-Offload with DeepSpeedCPUAdam. But current setting is ZeRO-Offload:{self.cpu_offload} and optimizer type {type(self.optimizer)}." \
286            f"Either disable fp16_master_weights_and_gradients or enable {self.zero_stage_string} Offload with DeepSpeedCPUAdam."
287
288        if self.reduce_scatter and self.partition_gradients:
289            valid_reduce_scatter_dtypes = (torch.float16, torch.bfloat16, torch.float32)
290            assert self.communication_data_type in valid_reduce_scatter_dtypes, f"{self.zero_stage_string} supports {valid_reduce_scatter_dtypes} communication_data_type with reduce scatter enabled. Got: '{self.communication_data_type}'"
291            assert self.gradient_predivide_factor == 1.0, f"gradient_predivide_factor != 1.0 is not yet supported with {self.zero_stage_string} with reduce scatter enabled"
292            assert self.postscale_gradients, f"pre-scale gradients is not yet supported with {self.zero_stage_string} with reduce scatter enabled"
293
294        # param flattened by groups
295        self.bit16_groups = []
296        self.bit16_groups_flat = []
297
298        # param partitioned by data parallel degree
299        # this will contain a list of equal sized tensors
300        # each of which will be updated by a different process
301        self.parallel_partitioned_bit16_groups = []
302
303        # a single 32-bit partition of the parallel partitioned parameters
304        # that this process will update
305        self.single_partition_of_fp32_groups = []
306
307        # param partition info
308
309        # These are the parameters in each group that will not be updated by this process directly
310        self.params_not_in_partition = []
311
312        # These are the parameters that will be updated by this process directly
313        self.params_in_partition = []
314
315        # Offset from the first parameter in the self.params_in_partition
316        # the parameter boundaries may not align with partition boundaries
317        # so we need to keep track of the offset
318        self.first_offset = []
319
320        # number of elements per partition in each group
321        self.partition_size = []
322
323        # align nccl all-gather send buffers to 4-byte boundary
324        self.nccl_start_alignment_factor = 2  # 4-byte alignment/sizeof(fp16) = 2
325
326        assert (
327            allgather_bucket_size % self.nccl_start_alignment_factor == 0
328        ), f"allgather_bucket_size must be a multiple of nccl_start_alignment_factor, {self.nccl_start_alignment_factor} "
329
330        self.all_reduce_print = False
331        self.dtype = self.optimizer.param_groups[0]['params'][0].dtype
332        self.gradient_accumulation_dtype = gradient_accumulation_dtype
333
334        if self.dtype != self.gradient_accumulation_dtype:
335            self.use_separate_grad_accum = True
336        else:
337            self.use_separate_grad_accum = False
338        if self.use_separate_grad_accum and not self.partition_gradients:
339            self.use_grad_accum_attribute = True
340        else:
341            self.use_grad_accum_attribute = False
342
343        self.round_robin_bit16_groups = []
344        self.round_robin_bit16_indices = []
345        self.round_robin_bit16_meta = []
346
347        # Use different parallel to do all_to_all_reduce related things
348        # padding on each partition for alignment purposes
349        self.groups_padding = []
350        # loop to deal with groups
351        for i, param_group in enumerate(self.optimizer.param_groups):
352            partition_id = dist.get_rank(group=self.real_dp_process_group[i])
353
354            # push this group to list before modify
355            # TODO: Explore simplification that avoids the extra book-keeping by pushing the reordered group
356            trainable_parameters = []
357            for param in param_group['params']:
358                if param.requires_grad:
359                    param.grad_accum = None
360                    trainable_parameters.append(param)
361            self.bit16_groups.append(trainable_parameters)
362
363            # not sure why apex was cloning the weights before flattening
364            # removing cloning here
365
366            see_memory_usage(f"Before moving param group {i} to CPU")
367            # move all the parameters to cpu to free up GPU space for creating flat buffer
368
369            # Create temp CPU param copies, free accelerator tensors
370            orig_group_numel = 0
371            for param in self.bit16_groups[i]:
372                orig_group_numel += param.numel()
373                param.cpu_data = opguard_wrapped__cpu_data_cpu(325)(param.data.cpu)
374                param.data = opguard_wrapped__data_to(326)(torch.empty(1).to, param.device)
375
376            empty_cache()
377            see_memory_usage(f"After moving param group {i} to CPU", force=False)
378
379            # Reorder group parameters for load balancing of gradient partitioning during backward among ranks.
380            # This ensures that gradients are reduced in a fashion such that ownership round robins among the ranks.
381            # For example, rather than 3 gradients (g_n+2, g_n+1, g_n) that are reduced consecutively belonging
382            # to the same rank, instead they will belong to 3 ranks (r_m+2, r_m+1, r_m).
383            if self.round_robin_gradients:
384                round_robin_tensors, round_robin_indices = self._round_robin_reorder(
385                    self.bit16_groups[i], dist.get_world_size(group=self.real_dp_process_group[i]))
386            else:
387                round_robin_tensors = self.bit16_groups[i]
388                round_robin_indices = list(range(len(self.bit16_groups[i])))
389
390            self.round_robin_bit16_groups.append(round_robin_tensors)
391            self.round_robin_bit16_indices.append(round_robin_indices)
392
393            # Create meta tensors list, ordered according to round_robin_tensors
394            meta_tensors = []
395            for param in round_robin_tensors:
396                opguard_wrapped__meta_tensors_append(348)(meta_tensors.append, torch.zeros_like(param.cpu_data, device="meta"))
397            self.round_robin_bit16_meta.append(meta_tensors)
398
399            # create flat buffer in CPU
400            flattened_buffer = self.flatten_dense_tensors_aligned(
401                self.round_robin_bit16_groups[i],
402                self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[i]),
403                use_cpu_data=True)
404
405            # free temp CPU params
406            for param in self.bit16_groups[i]:
407                del param.cpu_data
408
409            # Move CPU flat tensor to the accelerator memory.
410            opguard_wrapped__bit16_groups_flat_append(362)(self.bit16_groups_flat.append, flattened_buffer.to(get_accelerator().current_device_name()))
411            del flattened_buffer
412
413            see_memory_usage(f"After flattening and moving param group {i} to GPU", force=False)
414
415            # Record padding required for alignment
416            if partition_id == dist.get_world_size(group=self.real_dp_process_group[i]) - 1:
417                padding = self.bit16_groups_flat[i].numel() - orig_group_numel
418            else:
419                padding = 0
420            self.groups_padding.append(padding)
421
422            if dist.get_rank(group=self.real_dp_process_group[i]) == 0:
423                see_memory_usage(f"After Flattening and after emptying param group {i} cache", force=False)
424
425            # set model bit16 weight to slices of flattened buffer
426            self._update_model_bit16_weights(i)
427
428            # divide the flat weights into near equal partition equal to the data parallel degree
429            # each process will compute on a different part of the partition
430            data_parallel_partitions = self.get_data_parallel_partitions(self.bit16_groups_flat[i], i)
431            self.parallel_partitioned_bit16_groups.append(data_parallel_partitions)
432
433            # verify that data partition start locations are 4-byte aligned
434            for partitioned_data in data_parallel_partitions:
435                assert (partitioned_data.data_ptr() % (2 * self.nccl_start_alignment_factor) == 0)
436
437            # A partition of the fp32 master weights that will be updated by this process.
438            # Note that the params in single_partition_of_fp32_groups is cloned and detached
439            # from the origin params of the model.
440            if not fp16_master_weights_and_gradients:
441                weights_partition = opguard_wrapped__weights_partition_detach(393)(self.parallel_partitioned_bit16_groups[i][partition_id].to(
442                    self.device).clone().float().detach)
443            else:
444                weights_partition = opguard_wrapped__weights_partition_detach(396)(self.parallel_partitioned_bit16_groups[i][partition_id].to(
445                    self.device).clone().half().detach)
446
447            if self.cpu_offload:
448                weights_partition = get_accelerator().pin_memory(weights_partition)
449
450            self.single_partition_of_fp32_groups.append(weights_partition)
451
452            # Set local optimizer to have flat params of its own partition.
453            # After this, the local optimizer will only contain its own partition of params.
454            # In that case, the local optimizer only saves the states(momentum, variance, etc.) related to its partition's params(zero stage1).
455            self.single_partition_of_fp32_groups[
456                i].requires_grad = True  # keep this in case internal optimizer uses it
457            param_group['params'] = [self.single_partition_of_fp32_groups[i]]
458
459            partition_size = len(self.bit16_groups_flat[i]) / dist.get_world_size(group=self.real_dp_process_group[i])
460            params_in_partition, params_not_in_partition, first_offset = self.get_partition_info(
461                self.round_robin_bit16_groups[i], partition_size, partition_id)
462
463            self.partition_size.append(partition_size)
464            self.params_in_partition.append(params_in_partition)
465            self.params_not_in_partition.append(params_not_in_partition)
466            self.first_offset.append(first_offset)
467
468        self.reduce_bucket_size = int(reduce_bucket_size)
469        self.use_multi_rank_bucket_allreduce = use_multi_rank_bucket_allreduce
470        self.allgather_bucket_size = int(allgather_bucket_size)
471
472        self.reduction_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream()
473        #self.copy_grad_stream = get_accelerator().Stream()
474        self.callback_queued = False
475
476        self.param_dict = {}
477
478        # map between param_id and bool to specify if a param is in this partition
479        self.is_param_in_current_partition = {}
480
481        self.grads_in_ipg_bucket = []
482        self.params_in_ipg_bucket = []
483        self.elements_in_ipg_bucket = 0
484        self.params_already_reduced = []
485        self._release_ipg_buffers()
486        self.previous_reduced_grads = None
487        self.ipg_bucket_has_moe_params = False
488
489        # simplified param id
490        self.param_id = {}
491
492        #interesting code: unique ids being assigned to individual parameters
493        largest_param_numel = 0
494        count = 0
495        for i, params_group in enumerate(self.bit16_groups):
496            for param in params_group:
497                unique_id = id(param)
498                self.param_id[unique_id] = count
499                self.param_dict[count] = param
500                self.params_already_reduced.append(False)
501                if param.numel() > largest_param_numel:
502                    largest_param_numel = param.numel()
503                count = count + 1
504
505        for param_group in self.params_in_partition:
506            for param in param_group:
507                self.is_param_in_current_partition[self.get_param_id(param)] = True
508
509        for param_group in self.params_not_in_partition:
510            for param in param_group:
511                self.is_param_in_current_partition[self.get_param_id(param)] = False
512
513        if self.cpu_offload:
514            self.accumulated_grads_in_cpu = {}
515            self.norm_for_param_grads = {}
516            self.local_overflow = False
517            self.grad_position = {}
518            self.temp_grad_buffer_for_cpu_offload = torch.zeros(largest_param_numel,
519                                                                device=self.device,
520                                                                dtype=self.dtype)
521            if self.cpu_offload_pin_memory:
522                self.temp_grad_buffer_for_cpu_offload = get_accelerator().pin_memory(
523                    self.temp_grad_buffer_for_cpu_offload)
524            self.temp_grad_buffer_for_gpu_offload = torch.zeros(largest_param_numel,
525                                                                device=get_accelerator().current_device_name(),
526                                                                dtype=self.dtype)
527            for i, params_group in enumerate(self.bit16_groups):
528                self.get_grad_position(i, self.params_in_partition[i], self.first_offset[i], self.partition_size[i])
529
530        # mapping from parameter to partition that it belongs to
531        self.param_to_partition_ids = {}
532
533        # stores if a partition has been reduced in this step
534        self.is_partition_reduced = {}
535
536        # number of grads in partition that still need to be computed
537        self.remaining_grads_in_partition = {}
538
539        # total number of grads in partition
540        self.total_grads_in_partition = {}
541
542        # stores if a grad in a partition has been computed or not
543        self.is_grad_computed = {}
544
545        # stores the offset at which a parameter gradient needs to be inserted in a partition
546        self.grad_partition_insertion_offset = {}
547
548        # the offset in the gradient at which it must be inserted at the beginning of the partition
549        self.grad_start_offset = {}
550
551        # will store the averaged gradients required by this partition
552        self.averaged_gradients = {}
553
554        # For cpu_offload, will store the averaged gradients required by this partition
555        self.offload_gradient_dict = {}
556
557        # store index of first parameter in each partition
558        self.first_param_index_in_partition = {}
559
560        # initializes all data structures for implementing gradient partitioning
561        self.initialize_gradient_partitioning_data_structures()
562
563        # resets the data structure value for the next backward propagation
564        self.reset_partition_gradient_structures()
565
566        # creates backward hooks for gradient partitioning
567        self._grad_acc_hooks = []
568        if self.partition_gradients or self.overlap_comm:
569            self.create_reduce_and_remove_grad_hooks()
570
571        self.custom_loss_scaler = False
572        self.external_loss_scale = None
573
574        # we may have a way of fusing dynamic scale. Do not support for now
575        self.loss_scaler = CreateLossScaler(dtype=self.dtype,
576                                            static_loss_scale=static_loss_scale,
577                                            dynamic_scaling=dynamic_loss_scale,
578                                            dynamic_loss_args=dynamic_loss_args)
579        self.dynamic_loss_scale = self.loss_scaler.dynamic
580
581        if self.dtype != torch.float16:
582            # Only fp16 should use dynamic loss scaling
583            assert self.loss_scaler.cur_scale == 1.0
584            assert not self.dynamic_loss_scale
585
586        see_memory_usage("Before initializing optimizer states", force=True)
587        self.initialize_optimizer_states()
588        see_memory_usage("After initializing optimizer states", force=True)
589
590        if dist.get_rank() == 0:
591            logger.info(f"optimizer state initialized")
592
593        if dist.get_rank(group=self.dp_process_group) == 0:
594            see_memory_usage(f"After initializing ZeRO optimizer", force=True)
595
596        self._link_all_hp_params()
597        self._hp_optimizer_states_linked = False
598
599        self._enable_universal_checkpoint()
600        self._param_slice_mappings = self._create_param_mapping()
601
602    def destroy(self):
603        for hook in self._grad_acc_hooks:
604            hook.remove()
605        self.print_rank_0("Removed grad acc hooks")
606
607    def _enable_universal_checkpoint(self):
608        for lp_param_group in self.bit16_groups:
609            enable_universal_checkpoint(param_list=lp_param_group)
610
611    def _create_param_mapping(self):
612        param_mapping = []
613        for i, _ in enumerate(self.optimizer.param_groups):
614            param_mapping_per_group = OrderedDict()
615            for lp in self.bit16_groups[i]:
616                if lp._hp_mapping is not None:
617                    lp_name = self.param_names[lp]
618                    param_mapping_per_group[lp_name] = lp._hp_mapping.get_hp_fragment_address()
619            param_mapping.append(param_mapping_per_group)
620
621        return param_mapping
622
623    def _link_all_hp_params(self):
624        dp_world_size = dist.get_world_size(group=self.dp_process_group)
625        if self.cpu_offload:
626            self._get_offload_gradient_dict()
627
628        for i, _ in enumerate(self.optimizer.param_groups):
629            # Link bit16 and fp32 params in partition
630            partition_id = dist.get_rank(group=self.real_dp_process_group[i])
631            partition_size = self.bit16_groups_flat[i].numel() // dp_world_size
632            flat_hp_partition = self.single_partition_of_fp32_groups[i]
633            link_hp_params(lp_param_list=self.bit16_groups[i],
634                           flat_hp_partition=flat_hp_partition,
635                           gradient_dict=self.averaged_gradients,
636                           offload_gradient_dict=self.offload_gradient_dict,
637                           use_offload=self.cpu_offload,
638                           param_group_index=i,
639                           partition_start=partition_id * partition_size,
640                           partition_size=partition_size,
641                           dp_group=self.real_dp_process_group[i])
642
643    def _lazy_init_hp_params_optimizer_state(self):
644        if not self._hp_optimizer_states_linked:
645            for i, _ in enumerate(self.optimizer.param_groups):
646                lazy_init_hp_params_optimizer_state(self.bit16_groups[i], self.single_partition_of_fp32_groups[i],
647                                                    self.optimizer.state)
648            self._hp_optimizer_states_linked = True
649
650    def is_moe_group(self, group):
651        return 'moe' in group and group['moe']
652
653    def _configure_moe_settings(self):
654        # if we're using ZeRO stage 2, ensure contiguous gradients are used
655        if self.partition_gradients:
656            assert self.contiguous_gradients, "Contiguous Gradients in ZeRO Stage 2 must be set to True for MoE. Other code paths are not tested with MoE"
657        # NOTE: To run ZeRO stage 1 with MoE, we need to set self.contiguous_gradients to True or ignore the assertion
658        if not self.partition_gradients and not self.contiguous_gradients:
659            logger.warn(
660                "ZeRO Stage 1 has not been thoroughly tested with MoE. This configuration is still experimental.")
661        assert self.reduce_scatter, "Reduce Scatter in ZeRO Stage 2 must be set to True for MoE. Other code paths are not tested with MoE"
662
663        assert any(
664            [self.is_moe_group(group) for group in self.optimizer.param_groups]
665        ), "The model has moe layers, but None of the param groups are marked as MoE. Create a param group with 'moe' key set to True before creating optimizer"
666        self.is_moe_param_group = []
667        for i, group in enumerate(self.optimizer.param_groups):
668            if self.is_moe_group(group):
669                assert all([is_moe_param(param)
670                            for param in group['params']]), "All params in MoE group must be MoE params"
671                self.real_dp_process_group[i] = self.expert_dp_process_group[group['name']]
672                self.partition_count[i] = dist.get_world_size(group=self.expert_dp_process_group[group['name']])
673                self.is_moe_param_group.append(True)
674            else:
675                self.is_moe_param_group.append(False)
676
677        assert self.expert_dp_process_group is not None, "Expert data parallel group should be configured with MoE"
678        assert self.ep_process_group is not None, "Expert parallel group should be configured with MoE"
679
680    def _update_model_bit16_weights(self, group_index):
681        updated_params = opguard_wrapped__updated_params_unflatten(633)(self.unflatten, self.bit16_groups_flat[group_index], self.round_robin_bit16_meta[group_index])
682        for p, q in zip(self.round_robin_bit16_groups[group_index], updated_params):
683            p.data = q.data
684
685        # set model fp16 weight to slices of reordered flattened buffer
686        for param_index, param in enumerate(self.bit16_groups[group_index]):
687            new_index = self.round_robin_bit16_indices[group_index][param_index]
688            param.data = self.round_robin_bit16_groups[group_index][new_index].data
689
690    def _round_robin_reorder(self, tensor_list, num_partitions):
691
692        # disable round robin if need to debug something
693        # return tensor_list, list(range(len(tensor_list)))
694
695        partition_tensors = {}
696
697        for i, tensor in enumerate(tensor_list):
698            j = i % num_partitions
699            if not j in partition_tensors:
700                partition_tensors[j] = []
701            partition_tensors[j].append((i, tensor))
702
703        reordered_tensors = []
704        reordered_indices = {}
705
706        for partition_index in partition_tensors.keys():
707            for i, (original_index, tensor) in enumerate(partition_tensors[partition_index]):
708                reordered_indices[original_index] = len(reordered_tensors)
709                reordered_tensors.append(tensor)
710
711        return reordered_tensors, reordered_indices
712
713    def _release_ipg_buffers(self):
714        if self.contiguous_gradients:
715            self.ipg_buffer = None
716            self.grads_in_partition = None
717            self.grads_in_partition_offset = 0
718
719    def initialize_optimizer_states(self):
720
721        for i, group in enumerate(self.bit16_groups):
722            single_grad_partition = opguard_wrapped__single_grad_partition_zeros(674)(torch.zeros, int(self.partition_size[i]),
723                                                dtype=self.single_partition_of_fp32_groups[i].dtype,
724                                                device=self.device)
725            self.single_partition_of_fp32_groups[i].grad = get_accelerator().pin_memory(
726                single_grad_partition) if self.cpu_offload_pin_memory else single_grad_partition
727
728        # Initialize the optimizer states with the flattened fp32 partition.
729        # State initialization for the Adagrad optimizer occurs at construction as opposed to other optimizers
730        # which do lazy initialization of the state at the first call to step.
731        if isinstance(self.optimizer, torch.optim.Adagrad):
732            self.optimizer = torch.optim.Adagrad(self.single_partition_of_fp32_groups, **self.optimizer.defaults)
733
734        if not self.cpu_offload:
735            for group in self.single_partition_of_fp32_groups:
736                group.grad = None  #class init
737
738        return
739
740    #########################################################################
741    #################### ZeRO Stage 1 - reduce gradients ####################
742    #########################################################################
743    def reduce_gradients(self, pipeline_parallel=False):
744        world_size = dist.get_world_size(self.dp_process_group)
745        my_rank = dist.get_rank(self.dp_process_group)
746
747        # with PP we must create ipg buffer, since backward is handled outside zero
748        if pipeline_parallel and self.contiguous_gradients:
749            self.ipg_buffer = []
750            buf_0 = opguard_wrapped__buf_0_empty(702)(torch.empty, int(self.reduce_bucket_size),
751                                dtype=self.dtype,
752                                device=get_accelerator().current_device_name())
753            self.ipg_buffer.append(buf_0)
754            self.ipg_index = 0
755
756        if not self.overlap_comm:
757            for i, group in enumerate(self.bit16_groups):
758                for param in group:
759                    grad_reduc = self.get_gradient_for_reduction(param)
760                    if grad_reduc is not None:
761                        self.reduce_ready_partitions_and_remove_grads(param, i)
762        # reduce any pending grads in either hook/non-hook case
763        self.overlapping_partition_gradients_reduce_epilogue()
764
765    #########################################################################
766    #########################ZeRO Partition Gradients########################
767    #########################################################################
768
769    def get_first_param_index(self, group_id, param_group, partition_id):
770        for index, param in enumerate(param_group):
771            param_id = self.get_param_id(param)
772            if partition_id in self.param_to_partition_ids[group_id][param_id]:
773                return index
774        return None
775
776    def initialize_gradient_partitioning_data_structures(self):
777
778        for i, param_group in enumerate(self.round_robin_bit16_groups):
779            total_partitions = dist.get_world_size(group=self.real_dp_process_group[i])
780
781            self.param_to_partition_ids[i] = {}
782            self.is_partition_reduced[i] = {}
783            self.total_grads_in_partition[i] = {}
784            self.remaining_grads_in_partition[i] = {}
785            self.is_grad_computed[i] = {}
786            self.grad_partition_insertion_offset[i] = {}
787            self.grad_start_offset[i] = {}
788            self.first_param_index_in_partition[i] = {}
789
790            for partition_id in range(total_partitions):
791                self.is_grad_computed[i][partition_id] = {}
792                self.grad_partition_insertion_offset[i][partition_id] = {}
793                self.grad_start_offset[i][partition_id] = {}
794                self.total_grads_in_partition[i][partition_id] = 0
795                self.initialize_gradient_partition(i, param_group, partition_id)
796                self.is_partition_reduced[i][partition_id] = False
797                self.first_param_index_in_partition[i][partition_id] = self.get_first_param_index(
798                    i, param_group, partition_id)
799
800    def independent_gradient_partition_epilogue(self):
801        self.report_ipg_memory_usage(f"In ipg_epilogue before reduce_ipg_grads", 0)
802        self.reduce_ipg_grads()
803        self.report_ipg_memory_usage(f"In ipg_epilogue after reduce_ipg_grads", 0)
804
805        # if dist.get_rank() == 0:
806        #    logger.info("Params already reduced %s", self.params_already_reduced)
807        for i in range(len(self.params_already_reduced)):
808            self.params_already_reduced[i] = False
809
810        if self.overlap_comm:
811            if not get_accelerator().resolves_data_dependency():
812                get_accelerator().synchronize()
813            # It is safe to clear previously reduced grads of other partitions
814            self._clear_previous_reduced_grads()
815
816        if self.cpu_offload is False:
817            for i, _ in enumerate(self.bit16_groups):
818
819                if not i in self.averaged_gradients or self.averaged_gradients[i] is None:
820                    self.averaged_gradients[i] = self.get_flat_partition(
821                        self.params_in_partition[i],
822                        self.first_offset[i],
823                        self.partition_size[i],
824                        dtype=self.gradient_accumulation_dtype,
825                        device=get_accelerator().current_device_name(),
826                        return_tensor_list=True)
827                else:
828                    avg_new = self.get_flat_partition(self.params_in_partition[i],
829                                                      self.first_offset[i],
830                                                      self.partition_size[i],
831                                                      dtype=self.gradient_accumulation_dtype,
832                                                      device=get_accelerator().current_device_name(),
833                                                      return_tensor_list=True)
834
835                    for accumulated_grad, new_avg_grad in zip(self.averaged_gradients[i], avg_new):
836                        accumulated_grad.add_(new_avg_grad)
837
838        self._release_ipg_buffers()
839
840        # No need to keep the gradients anymore.
841        # All gradients required by the step
842        # are in self.averaged_gradients
843        self.zero_grad(set_to_none=True)
844        see_memory_usage(f"End ipg_epilogue")
845
846    # resets all partition to no reduced
847    # sets remaining grads to the total number of grads in each partition
848    # set is grad computed to false for all grads in partition
849    def reset_partition_gradient_structures(self):
850        for i, _ in enumerate(self.bit16_groups):
851            total_partitions = dist.get_world_size(group=self.real_dp_process_group[i])
852            for partition_id in range(total_partitions):
853                self.is_partition_reduced[i][partition_id] = False
854                self.remaining_grads_in_partition[i][partition_id] = self.total_grads_in_partition[i][partition_id]
855
856                for param_id in self.is_grad_computed[i][partition_id]:
857                    self.is_grad_computed[i][partition_id][param_id] = False
858
859    def initialize_gradient_partition(self, i, param_group, partition_id):
860
861        def set_key_value_list(dictionary, key, value):
862            if key in dictionary:
863                dictionary[key].append(value)
864            else:
865                dictionary[key] = [value]
866
867        def increment_value(dictionary, key):
868            if key in dictionary:
869                dictionary[key] += 1
870            else:
871                dictionary[key] = 1
872
873        partition_size = self.partition_size[i]
874
875        start_index = partition_size * partition_id
876        end_index = partition_size * (partition_id + 1)
877
878        current_index = 0
879        first_offset = 0
880
881        for param in param_group:
882
883            param_size = param.numel()
884            param_id = self.get_param_id(param)
885
886            if start_index <= current_index < end_index:
887                set_key_value_list(self.param_to_partition_ids[i], param_id, partition_id)
888                increment_value(self.total_grads_in_partition[i], partition_id)
889
890                self.is_grad_computed[i][partition_id][param_id] = False
891
892                self.grad_partition_insertion_offset[i][partition_id][param_id] = current_index - start_index
893                self.grad_start_offset[i][partition_id][param_id] = 0
894
895            elif current_index < start_index < (current_index + param_size):
896                assert (first_offset == 0
897                        ), "This can happen either zero or only once as this must be the first tensor in the partition"
898                first_offset = start_index - current_index
899
900                set_key_value_list(self.param_to_partition_ids[i], param_id, partition_id)
901                increment_value(self.total_grads_in_partition[i], partition_id)
902
903                self.is_grad_computed[i][partition_id][param_id] = False
904
905                self.grad_partition_insertion_offset[i][partition_id][param_id] = 0
906                self.grad_start_offset[i][partition_id][param_id] = first_offset
907
908            current_index = current_index + param_size
909
910    def overlapping_partition_gradients_reduce_epilogue(self):
911        self.independent_gradient_partition_epilogue()
912
913    def fill_grad_accum_attribute(self):
914        for group in self.bit16_groups:
915            for param in group:
916                if param.grad is not None:
917                    if param.grad_accum is None:
918                        param.grad_accum = param.grad.to(self.gradient_accumulation_dtype)
919                    else:
920                        param.grad_accum.add_(
921                            param.grad.to(self.gradient_accumulation_dtype).view(param.grad_accum.shape))
922                    param.grad = None
923
924    def get_gradient_for_reduction(self, param):
925        if self.use_grad_accum_attribute:
926            return param.grad_accum.to(self.dtype) if param.grad_accum is not None else None
927        else:
928            return param.grad
929
930    def get_param_gradient_attribute(self, param):
931        return param.grad_accum if self.use_grad_accum_attribute else param.grad
932
933    # Clear the tensor the reduction gradient attribute is pointing to
934    def clear_grad_attribute(self, param):
935        if self.use_grad_accum_attribute:
936            param.grad_accum = None
937        else:
938            param.grad = None
939
940    def create_reduce_and_remove_grad_hooks(self):
941        self.grad_accs = []
942        for i, param_group in enumerate(self.bit16_groups):
943            for param in param_group:
944                if param.requires_grad:
945
946                    def wrapper(param, i):
947                        param_tmp = param.expand_as(param)
948                        grad_acc = param_tmp.grad_fn.next_functions[0][0]
949
950                        def reduce_partition_and_remove_grads(*notneeded):
951                            self.reduce_ready_partitions_and_remove_grads(param, i)
952
953                        self._grad_acc_hooks.append(grad_acc.register_hook(reduce_partition_and_remove_grads))
954                        self.grad_accs.append(grad_acc)
955
956                    opguard_wrapped__param_wrapper(908)(wrapper, param, i)
957
958    def get_param_id(self, param):
959        unique_id = id(param)
960        return self.param_id[unique_id]
961
962    def report_ipg_memory_usage(self, tag, param_elems):
963        elem_count = self.elements_in_ipg_bucket + param_elems
964        percent_of_bucket_size = (100.0 * elem_count) // self.reduce_bucket_size
965        see_memory_usage(
966            f"{tag}: elems in_bucket {self.elements_in_ipg_bucket} param {param_elems} max_percent {percent_of_bucket_size}"
967        )
968
969    # create a flat tensor aligned at the alignment boundary
970    def flatten_dense_tensors_aligned(self, tensor_list, alignment, use_cpu_data=False):
971        tensor_list = [param.cpu_data for param in tensor_list] if use_cpu_data else tensor_list
972        return opguard_wrapped__align_dense_tensors_flatten(924)(self.flatten, align_dense_tensors(tensor_list, alignment))
973
974    ############### Independent Partition Gradient ########################
975    def reduce_independent_p_g_buckets_and_remove_grads(self, param, i):
976
977        grad_reduc = self.get_gradient_for_reduction(param)
978        if self.elements_in_ipg_bucket + param.numel() > self.reduce_bucket_size:
979            self.report_ipg_memory_usage("In ipg_remove_grads before reduce_ipg_grads", param.numel())
980            self.reduce_ipg_grads()
981            if self.contiguous_gradients and self.overlap_comm:
982                # Swap ipg_index between 0 and 1
983                self.ipg_index = 1 - self.ipg_index
984            self.report_ipg_memory_usage("In ipg_remove_grads after reduce_ipg_grads", param.numel())
985
986        param_id = self.get_param_id(param)
987        assert self.params_already_reduced[param_id] == False, \
988            f"The parameter {param_id} has already been reduced. \
989            Gradient computed twice for this partition. \
990            Multiple gradient reduction is currently not supported"
991
992        if self.contiguous_gradients:
993            if param.numel() > self.reduce_bucket_size:
994                self.extra_large_param_to_reduce = param
995            else:
996                # keeping the gradients contiguous to prevent memory fragmentation, and avoid flattening
997                new_grad_tensor = opguard_wrapped__new_grad_tensor_narrow(949)(self.ipg_buffer[self.ipg_index].narrow, 0, self.elements_in_ipg_bucket, param.numel())
998                opguard_wrapped__new_grad_tensor_copy_(950)(new_grad_tensor.copy_, grad_reduc.view(-1))
999                grad_reduc.data = opguard_wrapped__data_view_as(951)(new_grad_tensor.data.view_as, grad_reduc)
1000
1001        self.elements_in_ipg_bucket += param.numel()
1002
1003        assert grad_reduc is not None, f"rank {dist.get_rank()} - Invalid to reduce Param {param_id} with None gradient"
1004
1005        self.grads_in_ipg_bucket.append(grad_reduc)
1006        self.params_in_ipg_bucket.append((i, param, param_id))
1007
1008        #make sure the average tensor function knows how to average the gradients
1009        if is_moe_param(param):
1010            self.ipg_bucket_has_moe_params = True
1011
1012        self.report_ipg_memory_usage("End ipg_remove_grads", 0)
1013
1014    def print_rank_0(self, message):
1015        if dist.get_rank() == 0:
1016            logger.info(message)
1017
1018    def gradient_reduction_w_predivide(self, tensor):
1019
1020        dp_world_size = dist.get_world_size(group=self.dp_process_group)
1021
1022        tensor_to_allreduce = tensor
1023
1024        if self.communication_data_type != tensor.dtype:
1025            tensor_to_allreduce = tensor.to(self.communication_data_type)
1026
1027        if self.postscale_gradients:
1028            if self.gradient_predivide_factor != 1.0:
1029                tensor_to_allreduce.mul_(1. / self.gradient_predivide_factor)
1030
1031            dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group)
1032
1033            if self.gradient_predivide_factor != dp_world_size:
1034                tensor_to_allreduce.mul_(self.gradient_predivide_factor /
1035                                         (dp_world_size / float(self.sequence_parallel_size)))
1036        else:
1037            tensor_to_allreduce.div_(dp_world_size / float(self.sequence_parallel_size))
1038            dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group)
1039
1040        if self.communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce:
1041            tensor.copy_(tensor_to_allreduce)
1042
1043        return tensor
1044
1045    def allreduce_and_copy_with_multiple_ranks(self,
1046                                               small_bucket,
1047                                               log=None,
1048                                               divide=True,
1049                                               process_group=None,
1050                                               bucket_ranks=None):
1051        process_group = self.dp_process_group if process_group is None else process_group
1052        allreduced = self.allreduce_bucket(small_bucket, log=log, divide=divide, process_group=process_group)
1053        for buf, synced, bucket_rank in zip(small_bucket, opguard_wrapped__allreduced_unflatten(1005)(self.unflatten, allreduced, small_bucket), bucket_ranks):
1054            if dist.get_rank(group=process_group) == bucket_rank:
1055                opguard_wrapped__buf_copy_(1007)(buf.copy_, synced)
1056
1057    def allreduce_and_scatter(self, bucket, numel_per_bucket=500000000, log=None, divide=True, process_group=None):
1058        small_bucket = []
1059        small_bucket_ranks = []
1060        numel = 0
1061        allreduce_sizes = []
1062
1063        for i, bucket_elem in enumerate(bucket):
1064            rank, tensor = bucket_elem
1065            small_bucket.append(tensor)
1066            small_bucket_ranks.append(rank)
1067            numel = numel + tensor.numel()
1068            if numel > numel_per_bucket:
1069                self.allreduce_and_copy_with_multiple_ranks(small_bucket,
1070                                                            log=None,
1071                                                            divide=divide,
1072                                                            process_group=process_group,
1073                                                            bucket_ranks=small_bucket_ranks)
1074                small_bucket = []
1075                small_bucket_ranks = []
1076                numel = 0
1077
1078        if len(small_bucket) > 0:
1079            self.allreduce_and_copy_with_multiple_ranks(small_bucket,
1080                                                        log=None,
1081                                                        divide=divide,
1082                                                        process_group=process_group,
1083                                                        bucket_ranks=small_bucket_ranks)
1084
1085    def average_tensor(self, tensor):
1086        if self.overlap_comm:
1087            stream = self.reduction_stream
1088            if not get_accelerator().resolves_data_dependency():
1089                stream.wait_stream(get_accelerator().current_stream())
1090        else:
1091            stream = get_accelerator().current_stream()
1092
1093        with get_accelerator().stream(stream):
1094            if not self.reduce_scatter:
1095                self.gradient_reduction_w_predivide(tensor)
1096                return
1097
1098            # Accumulate destination ranks and bucket offsets for each gradient slice.
1099            # Note: potential future optimization, record access pattern of parameters
1100            # in backward pass and partition gradients w.r.t. access pattern so that our
1101            # bucket is guaranteed to be contiguous w.r.t. ranks
1102            rank_and_offsets = []
1103            real_dp_process_group = []
1104            curr_size = 0
1105            prev_id, prev_process_group = -1, None
1106
1107            process_group = self.dp_process_group
1108            # count = 0
1109            for i, param, param_id in self.params_in_ipg_bucket:
1110
1111                process_group = self.dp_process_group
1112                grad_reduc = self.get_gradient_for_reduction(param)
1113                #Averages gradients at parameter level if ipg has a moe param
1114                #Otherwise averaging is done at the entire buffer level at the end of the loop
1115                # MoE param have different groups
1116                if self.ipg_bucket_has_moe_params:
1117                    process_group = self.expert_dp_process_group[param.group_name] if is_moe_param(
1118                        param) else self.dp_process_group
1119                    grad_reduc.data.div_(dist.get_world_size(group=process_group) / float(self.sequence_parallel_size))
1120
1121                partition_ids = self.param_to_partition_ids[i][param_id]
1122                assert all([p_id < dist.get_world_size(group=process_group) for p_id in partition_ids
1123                            ]), f"world size {dist.get_world_size(group=process_group)} and p_ids: {partition_ids}"
1124                partition_size = self.partition_size[i]
1125                # Get all partition ids + their offsets
1126                partition_ids_w_offsets = []
1127                for partition_id in partition_ids:
1128                    offset = self.grad_start_offset[i][partition_id][param_id]
1129                    partition_ids_w_offsets.append((partition_id, offset))
1130                partition_ids_w_offsets.sort(key=lambda t: t[1])
1131
1132                # Calculate rank and offsets for grad slices
1133                for idx in range(len(partition_ids_w_offsets)):
1134                    partition_id, offset = partition_ids_w_offsets[idx]
1135
1136                    # if dist.get_rank() == 0 and count < 100:
1137                    #     print(f"Rank {dist.get_rank()} rank offset id {idx} calculated dp size {dist.get_world_size(group=process_group)} real dp size {dist.get_world_size(self.real_dp_process_group[i])} and dst: {partition_id}")
1138                    # count += 1
1139
1140                    # Calculate numel for grad slice depending on partition location
1141                    if idx == len(partition_ids_w_offsets) - 1:
1142                        # Last partition_id uses its own offset
1143                        numel = param.numel() - offset
1144                    else:
1145                        # Set numel to next partition's offset
1146                        numel = partition_ids_w_offsets[idx + 1][1] - offset
1147
1148                    # Merge bucket ranges if they belong to the same rank
1149                    if partition_id == prev_id and process_group == prev_process_group:
1150                        prev_pid, prev_size, prev_numel = rank_and_offsets[-1]
1151                        rank_and_offsets[-1] = (prev_pid, prev_size, prev_numel + numel)
1152                    else:
1153                        rank_and_offsets.append((partition_id, curr_size, numel))
1154                        real_dp_process_group.append(process_group)
1155                    curr_size += numel
1156                    prev_id, prev_process_group = partition_id, process_group
1157
1158            if not self.ipg_bucket_has_moe_params:
1159                opguard_wrapped__tensor_div_(1111)(tensor.div_, dist.get_world_size(group=self.dp_process_group) / float(self.sequence_parallel_size))
1160
1161            buckets = {}
1162            for i, (dst, bucket_offset, numel) in enumerate(rank_and_offsets):
1163                grad_slice = opguard_wrapped__grad_slice_narrow(1115)(tensor.narrow, 0, int(bucket_offset), int(numel))
1164                bucket_key = real_dp_process_group[i] if self.use_multi_rank_bucket_allreduce else (
1165                    dst, real_dp_process_group[i])
1166                if bucket_key not in buckets:
1167                    buckets[bucket_key] = []
1168                if self.use_multi_rank_bucket_allreduce:
1169                    buckets[bucket_key].append((dst, grad_slice))
1170                else:
1171                    buckets[bucket_key].append(grad_slice)
1172
1173            for bucket_key in buckets:
1174                if self.use_multi_rank_bucket_allreduce:
1175                    self.allreduce_and_scatter(buckets[bucket_key],
1176                                               numel_per_bucket=self.reduce_bucket_size,
1177                                               divide=False,
1178                                               process_group=bucket_key)
1179                else:
1180                    dst, process_group = bucket_key
1181                    self.allreduce_no_retain(buckets[bucket_key],
1182                                             numel_per_bucket=self.reduce_bucket_size,
1183                                             rank=dst,
1184                                             divide=False,
1185                                             process_group=process_group)
1186
1187    ##############################################################################
1188    ############################# CPU Offload Methods#############################
1189    ##############################################################################
1190    def get_grad_position(self, group_id, tensor_list, first_offset, partition_size):
1191        current_offset = 0
1192
1193        for i, tensor in enumerate(tensor_list):
1194            param_id = self.get_param_id(tensor)
1195            param_start_offset = 0
1196
1197            num_elements = tensor.numel()
1198
1199            # we need to offset to get to the right element
1200            if i == 0 and first_offset > 0:
1201                tensor_offset = first_offset
1202                num_elements = num_elements - tensor_offset
1203                param_start_offset = first_offset
1204
1205            # we dont need all elements of the tensor
1206            if num_elements > (partition_size - current_offset):
1207                num_elements = partition_size - current_offset
1208
1209            self.grad_position[param_id] = [
1210                int(group_id), int(param_start_offset),
1211                int(current_offset), int(num_elements)
1212            ]
1213            current_offset += num_elements
1214
1215    def update_overflow_tracker_for_param_grad(self, param):
1216        grad_accum = self.get_param_gradient_attribute(param)
1217        if grad_accum is not None and self._has_inf_or_nan(grad_accum.data):
1218            self.local_overflow = True
1219
1220    def _get_offload_gradient_dict(self):
1221        for param_group_index, _ in enumerate(self.optimizer.param_groups):
1222            self.offload_gradient_dict[param_group_index] = []
1223            for lp_param in self.params_in_partition[param_group_index]:
1224                param_id = self.get_param_id(lp_param)
1225                [_, _, dest_offset, num_elements] = self.grad_position[param_id]
1226                dest_tensor = self.single_partition_of_fp32_groups[param_group_index].grad.view(-1).narrow(
1227                    0, dest_offset, num_elements)
1228                self.offload_gradient_dict[param_group_index].append(dest_tensor)
1229
1230    def async_accumulate_grad_in_cpu_via_gpu(self, param):
1231        param_id = self.get_param_id(param)
1232
1233        [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id]
1234
1235        # copy to a preexisiting buffer to avoid memory allocation penalty
1236        dest_buffer = self.temp_grad_buffer_for_gpu_offload.view(-1).narrow(0, 0, param.numel())
1237
1238        #buffer for storing gradients for this parameter in CPU
1239        def buffer_to_accumulate_to_in_cpu():
1240            if not self.fp16_master_weights_and_gradients:
1241                buffer = torch.zeros(param.numel(), dtype=param.dtype, device=self.device)
1242                return get_accelerator().pin_memory(buffer) if self.cpu_offload_pin_memory else buffer
1243            else:
1244                return self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements)
1245
1246        #accumulate gradients into param.grad_accum or parts of it that belongs to this partition
1247        def accumulate_gradients():
1248            grad_accum = self.get_param_gradient_attribute(param)
1249            if not self.fp16_master_weights_and_gradients:
1250                dest_buffer.copy_(self.accumulated_grads_in_cpu[param_id].view(-1), non_blocking=True)
1251                grad_accum.data.view(-1).add_(dest_buffer)
1252            else:
1253                dest_buffer.narrow(0, source_offset,
1254                                   num_elements).copy_(self.accumulated_grads_in_cpu[param_id].view(-1),
1255                                                       non_blocking=True)
1256                grad_accum.data.view(-1).narrow(0, source_offset,
1257                                                num_elements).add_(dest_buffer.narrow(0, source_offset, num_elements))
1258
1259        #move accumulated gradients back to CPU
1260        def copy_gradients_to_cpu():
1261            grad_accum = self.get_param_gradient_attribute(param)
1262            if not self.fp16_master_weights_and_gradients:
1263                self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1), non_blocking=True)
1264            else:
1265                self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1).narrow(
1266                    0, source_offset, num_elements),
1267                                                                   non_blocking=True)
1268
1269        if param_id not in self.accumulated_grads_in_cpu:
1270            self.accumulated_grads_in_cpu[param_id] = buffer_to_accumulate_to_in_cpu()
1271
1272        if self.micro_step_id > 0:
1273            accumulate_gradients()
1274
1275        # at the boundary we will send 32bit directly
1276        if not self.is_gradient_accumulation_boundary:
1277            copy_gradients_to_cpu()
1278
1279    def set_norm_for_param_grad(self, param):
1280        param_id = self.get_param_id(param)
1281        grad_accum = self.get_param_gradient_attribute(param)
1282        accumulated_grad = self.accumulated_grads_in_cpu[
1283            param_id] if self.gradient_accumulation_steps > 1 else grad_accum
1284
1285        [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id]
1286
1287        start = source_offset
1288        accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements)
1289
1290        self.norm_for_param_grads[param_id] = accumulated_grad.data.double().norm(2)
1291
1292    def set_norm_for_param_grad_in_gpu(self, param):
1293        param_id = self.get_param_id(param)
1294        grad_accum = self.get_param_gradient_attribute(param)
1295        if grad_accum is None:
1296            accumulated_grad = param.grad
1297        else:
1298            accumulated_grad = grad_accum
1299
1300        [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id]
1301
1302        start = source_offset
1303        accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements)
1304
1305        self.norm_for_param_grads[param_id] = accumulated_grad.data.double().norm(2)
1306
1307    def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param):
1308        param_id = self.get_param_id(param)
1309
1310        [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id]
1311
1312        dest_tensor = self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements)
1313
1314        grad_accum = self.get_param_gradient_attribute(param)
1315        if grad_accum is None:
1316            src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements)
1317        else:
1318            src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements)
1319        if not self.fp16_master_weights_and_gradients:
1320            src_tensor = src_tensor.float()
1321
1322        dest_tensor.copy_(src_tensor, non_blocking=True)
1323        param.grad = None  #offload only
1324
1325    def complete_grad_norm_calculation_for_cpu_offload(self, params):
1326        total_norm = 0.0
1327        norm_type = 2.0
1328        for p in params:
1329            # Pipeline parallelism may replicate parameters. Avoid multi-counting.
1330            if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated:
1331                continue
1332
1333            if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0):
1334                param_id = self.get_param_id(p)
1335                # as some model have trainable parameters but skipped in training,
1336                # their backward hooks in self.create_reduce_and_remove_grad_hooks() will not run,
1337                # so they have no norm_for_param_grads
1338                if param_id in self.norm_for_param_grads:
1339                    param_norm = self.norm_for_param_grads[param_id]
1340                    total_norm += param_norm.item()**2
1341                else:
1342                    # As unused parameters in modules may not be expected sometimes,
1343                    # add an explicit error msg when it occurred and an option to
1344                    # avoid the error
1345                    assert self.ignore_unused_parameters, """
1346                        This assert indicates that your module has parameters that
1347                        were not used in producing loss.
1348                        You can avoid this assert by
1349                        (1) enable ignore_unused_parameters option in zero_optimization config;
1350                        (2) making sure all trainable parameters and `forward` function
1351                            outputs participate in calculating loss.
1352                    """
1353
1354        # Sum across all model parallel GPUs.
1355        total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
1356        dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=self.dp_process_group)
1357
1358        self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.SUM)
1359
1360        total_norm = total_norm_cuda[0].item()**(1. / norm_type)
1361
1362        if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm:
1363            total_norm = -1
1364
1365        return total_norm
1366
1367    ############################################################################################
1368    def copy_grads_in_partition(self, param):
1369        if self.cpu_offload:
1370
1371            if self.gradient_accumulation_steps > 1:
1372                self.async_accumulate_grad_in_cpu_via_gpu(param)
1373
1374            if self.is_gradient_accumulation_boundary:
1375                self.set_norm_for_param_grad_in_gpu(param)
1376
1377                self.update_overflow_tracker_for_param_grad(param)
1378
1379                self.async_inplace_copy_grad_to_fp32_buffer_from_gpu(param)
1380
1381            return
1382        #print(f"ID {self.get_param_id(param)} grad norm {param.grad.norm()}")
1383        if self.grads_in_partition is None:
1384            self.grads_in_partition_offset = 0
1385            total_size = 0
1386            for group in self.params_in_partition:
1387                for param_in_partition in group:
1388                    total_size += param_in_partition.numel()
1389
1390            see_memory_usage(f"before copying {total_size} gradients into partition")
1391            self.grads_in_partition = opguard_wrapped__grads_in_partition_empty(1343)(torch.empty, int(total_size),
1392                                                  dtype=self.dtype,
1393                                                  device=get_accelerator().current_device_name())
1394            see_memory_usage(f"after copying {total_size} gradients into partition")
1395
1396        grad_reduc = self.get_gradient_for_reduction(param)
1397        # The allreduce buffer will be rewritten. Copy the gradients in partition to a new buffer
1398        new_grad_tensor = opguard_wrapped__new_grad_tensor_narrow(1350)(self.grads_in_partition.view(-1).narrow, 0, self.grads_in_partition_offset, param.numel())
1399        opguard_wrapped__new_grad_tensor_copy_(1351)(new_grad_tensor.copy_, grad_reduc.view(-1))
1400        grad_reduc.data = opguard_wrapped__data_view_as(1352)(new_grad_tensor.data.view_as, grad_reduc)
1401        #print(f"Grad norm after copy to contiguous_buffer {param.grad.data.norm()}")
1402        self.grads_in_partition_offset += param.numel()
1403
1404    def reduce_ipg_grads(self):
1405        if self.contiguous_gradients:
1406            if self.extra_large_param_to_reduce is not None:
1407                assert len(self.params_in_ipg_bucket) == 1, "more than 1 param in ipg bucket, this shouldn't happen"
1408                _, _, param_id = self.params_in_ipg_bucket[0]
1409                assert self.get_param_id(self.extra_large_param_to_reduce
1410                                         ) == param_id, "param in ipg bucket does not match extra-large param"
1411                extra_large_grad_reduc = self.get_gradient_for_reduction(self.extra_large_param_to_reduce)
1412                self.average_tensor(extra_large_grad_reduc.view(-1))
1413                self.extra_large_param_to_reduce = None
1414            else:
1415                opguard_wrapped__narrow_average_tensor(1367)(self.average_tensor, self.ipg_buffer[self.ipg_index].narrow(0, 0, self.elements_in_ipg_bucket))
1416        else:
1417            self.buffered_reduce_fallback(None,
1418                                          self.grads_in_ipg_bucket,
1419                                          elements_per_buffer=self.elements_in_ipg_bucket)
1420
1421        if self.overlap_comm:
1422            stream = self.reduction_stream
1423        elif self.cpu_offload:
1424            # TODO: copy_grad_stream is disabled because of race with reduce. This hurts perf and should be fixed.
1425            #            get_accelerator().synchronize()
1426            #            stream = self.copy_grad_stream
1427            stream = get_accelerator().current_stream()
1428        else:
1429            stream = get_accelerator().current_stream()
1430
1431        with get_accelerator().stream(stream):
1432            for _, param, param_id in self.params_in_ipg_bucket:
1433
1434                assert self.params_already_reduced[param_id] == False, \
1435                    f"The parameter {param_id} has already been reduced. \
1436                    Gradient computed twice for this partition. \
1437                    Multiple gradient reduction is currently not supported"
1438
1439                self.params_already_reduced[param_id] = True
1440                if self.partition_gradients:
1441                    if not self.is_param_in_current_partition[param_id]:
1442                        if self.overlap_comm and self.contiguous_gradients is False:
1443                            # Clear grads of other partitions during the next reduction
1444                            # to avoid clearing them before the reduction is complete.
1445                            if self.previous_reduced_grads is None:
1446                                self.previous_reduced_grads = []
1447                            self.previous_reduced_grads.append(param)
1448                        else:
1449                            self.clear_grad_attribute(param)
1450                    elif self.contiguous_gradients:
1451                        self.copy_grads_in_partition(param)
1452                else:  # zero stage 1 - partition only optimizer state
1453                    if self.contiguous_gradients and self.is_param_in_current_partition[param_id]:
1454                        self.copy_grads_in_partition(param)
1455
1456        self.grads_in_ipg_bucket = []
1457        self.params_in_ipg_bucket = []
1458        self.ipg_bucket_has_moe_params = False
1459        self.elements_in_ipg_bucket = 0
1460        #####################################################################
1461
1462    def reduce_ready_partitions_and_remove_grads(self, param, i):
1463        if self.partition_gradients or self.is_gradient_accumulation_boundary:
1464            self.reduce_independent_p_g_buckets_and_remove_grads(param, i)
1465
1466    def zero_reduced_gradients(self, partition_id, i):
1467
1468        def are_all_related_partitions_reduced(params_id):
1469            for partition_id in self.param_to_partition_ids[i][params_id]:
1470                if not self.is_partition_reduced[i][partition_id]:
1471                    return False
1472            return True
1473
1474        for params_id in self.is_grad_computed[i][partition_id]:
1475            if are_all_related_partitions_reduced(params_id):
1476                self.param_dict[params_id].grad = None  # dead code
1477
1478    def flatten_and_print(self, message, tensors, start=0, n=5):
1479        flatten_tensor = self.flatten(tensors)
1480
1481        def print_func():
1482            logger.info(flatten_tensor.contiguous().view(-1).narrow(0, start, n))
1483
1484        self.sequential_execution(print_func, message)
1485
1486    def get_grads_to_reduce(self, i, partition_id):
1487
1488        def get_reducible_portion(key):
1489            grad = self.param_dict[key].grad
1490            total_elements = grad.numel()
1491            start = self.grad_start_offset[i][partition_id][key]
1492            num_elements = min(total_elements - start,
1493                               self.partition_size[i] - self.grad_partition_insertion_offset[i][partition_id][key])
1494            if not pg_correctness_test:
1495                if num_elements == total_elements:
1496                    return grad
1497                else:
1498                    return grad.contiguous().view(-1).narrow(0, int(start), int(num_elements))
1499            else:
1500                if num_elements == total_elements:
1501                    return grad.clone()
1502                else:
1503                    return grad.clone().contiguous().view(-1).narrow(0, int(start), int(num_elements))
1504
1505        grads_to_reduce = []
1506        for key in self.is_grad_computed[i][partition_id]:
1507            grad = get_reducible_portion(key)
1508            grads_to_reduce.append(grad)
1509        return grads_to_reduce
1510
1511    def sequential_execution(self, function, message, group=None):
1512        if group is None:
1513            group = self.dp_process_group
1514        if dist.get_rank(group=group) == 0:
1515            logger.info(message)
1516        for id in range(dist.get_world_size(group=group)):
1517            if id == dist.get_rank(group=group):
1518                function()
1519            dist.barrier(group=group)
1520
1521    def set_none_gradients_to_zero(self, i, partition_id):
1522        for param_id in self.is_grad_computed[i][partition_id]:
1523            param = self.param_dict[param_id]
1524            if param.grad is None:
1525                param.grad = torch.zero_like(param)
1526
1527    ######################Reduction Related Methods##############################
1528    def allreduce_bucket(self, bucket, rank=None, log=None, divide=True, process_group=None):
1529        tensor = opguard_wrapped__tensor_flatten(1481)(self.flatten, bucket)
1530
1531        process_group = self.dp_process_group if process_group is None else process_group
1532
1533        tensor_to_allreduce = tensor
1534
1535        if pg_correctness_test or self.sequence_parallel_size > 1:
1536            communication_data_type = torch.float32
1537        else:
1538            communication_data_type = self.communication_data_type
1539
1540        if communication_data_type != tensor.dtype:
1541            tensor_to_allreduce = tensor.to(communication_data_type)
1542
1543        if divide:
1544            tensor_to_allreduce.div_(dist.get_world_size(group=process_group) / float(self.sequence_parallel_size))
1545
1546        if rank is None:
1547            #    "All Reducing"
1548            dist.all_reduce(tensor_to_allreduce, group=process_group)
1549        else:
1550            global_rank = dist.get_global_rank(process_group, rank)
1551            dist.reduce(tensor_to_allreduce, global_rank, group=process_group)
1552
1553        if communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce:
1554            if rank is None or rank == dist.get_rank(group=process_group):
1555                tensor.copy_(tensor_to_allreduce)
1556
1557        return tensor
1558
1559    def _clear_previous_reduced_grads(self):
1560        if self.previous_reduced_grads is not None:
1561            for param in self.previous_reduced_grads:
1562                self.clear_grad_attribute(param)
1563            self.previous_reduced_grads = None
1564
1565    # if rank is specified do a reduction instead of an allreduce
1566    def allreduce_and_copy(self, small_bucket, rank=None, log=None, divide=True, process_group=None):
1567        process_group = self.dp_process_group if process_group is None else process_group
1568        if self.overlap_comm:
1569            if not get_accelerator().resolves_data_dependency():
1570                get_accelerator().synchronize()
1571            # It is safe to clear the previously reduced grads of other partitions
1572            self._clear_previous_reduced_grads()
1573            stream = self.reduction_stream
1574        else:
1575            stream = get_accelerator().current_stream()
1576
1577        with get_accelerator().stream(stream):
1578            allreduced = self.allreduce_bucket(
1579                small_bucket,
1580                rank=rank,
1581                log=log,
1582                divide=divide,
1583                process_group=process_group,
1584            )
1585            if rank is None or rank == dist.get_rank(group=self.dp_process_group):
1586                for buf, synced in zip(small_bucket, opguard_wrapped__allreduced_unflatten(1538)(self.unflatten, allreduced, small_bucket)):
1587                    opguard_wrapped__buf_copy_(1539)(buf.copy_, synced)
1588
1589    def allreduce_no_retain(
1590        self,
1591        bucket,
1592        numel_per_bucket=500000000,
1593        rank=None,
1594        log=None,
1595        divide=True,
1596        process_group=None,
1597    ):
1598        small_bucket = []
1599        numel = 0
1600        for tensor in bucket:
1601            small_bucket.append(tensor)
1602            numel = numel + tensor.numel()
1603            if numel > numel_per_bucket:
1604                self.allreduce_and_copy(small_bucket, rank=rank, log=None, divide=divide, process_group=process_group)
1605                small_bucket = []
1606                numel = 0
1607
1608        if len(small_bucket) > 0:
1609            self.allreduce_and_copy(small_bucket, rank=rank, log=log, divide=divide, process_group=process_group)
1610
1611    # allows using reduction of gradients instead of using all_reduce
1612
1613    def buffered_reduce_fallback(self, rank, grads, elements_per_buffer=500000000, log=None):
1614        split_buckets = split_half_float_double(grads)
1615
1616        for i, bucket in enumerate(split_buckets):
1617            self.allreduce_no_retain(bucket, numel_per_bucket=elements_per_buffer, rank=rank, log=log)
1618
1619    #############################################################################
1620    #############################################################################
1621    #############################################################################
1622
1623    # views the tensor as multiple partitions and returns
1624    # those partitions
1625    def get_data_parallel_partitions(self, tensor, group_id):
1626        partitions = []
1627
1628        dp = dist.get_world_size(group=self.real_dp_process_group[group_id])
1629        # dp_id = dist.get_rank(group=self.real_dp_process_group[group_id])
1630
1631        total_num_elements = tensor.numel()
1632
1633        base_size = total_num_elements // dp
1634        remaining = total_num_elements % dp
1635
1636        start = 0
1637        for id in range(dp):
1638            partition_size = base_size
1639            if id < remaining:
1640                partition_size = partition_size + 1
1641            opguard_wrapped__partitions_append(1593)(partitions.append, tensor.narrow(0, start, partition_size))
1642            start = start + partition_size
1643        return partitions
1644
1645    def get_partition_info(self, tensor_list, partition_size, partition_id):
1646        params_in_partition = []
1647        params_not_in_partition = []
1648
1649        start_index = partition_size * partition_id
1650        end_index = partition_size * (partition_id + 1)
1651
1652        current_index = 0
1653        first_offset = 0
1654
1655        for tensor in tensor_list:
1656
1657            tensor_size = tensor.numel()
1658
1659            if start_index <= current_index < end_index:
1660                params_in_partition.append(tensor)
1661
1662            elif current_index < start_index < (current_index + tensor_size):
1663                params_in_partition.append(tensor)
1664
1665                assert (first_offset == 0
1666                        ), "This can happen either zero or only once as this must be the first tensor in the partition"
1667                first_offset = start_index - current_index
1668
1669            else:
1670                params_not_in_partition.append(tensor)
1671
1672            current_index = current_index + tensor_size
1673
1674        return params_in_partition, params_not_in_partition, first_offset
1675
1676    def zero_grad(self, set_to_none=True):
1677        """
1678        Zero FP16 parameter grads.
1679        """
1680        # FP32 grad should never exist.
1681        # For speed, set model fp16 grad to None by default
1682        # zero all pointers to grad tensors
1683        for group in self.bit16_groups:
1684            for p in group:
1685                if set_to_none:
1686                    p.grad = None  # epilogue and in step
1687                    p.grad_accum = None
1688                else:
1689                    if p.grad is not None:
1690                        p.grad.detach_()
1691                        p.grad.zero_()
1692
1693    def _model_parallel_all_reduce(self, tensor, op):
1694        """ Perform all reduce within model parallel group, if any.
1695        """
1696        if self.model_parallel_group is None or self.model_parallel_world_size == 1:
1697            pass
1698        else:
1699            dist.all_reduce(tensor=tensor, op=op, group=self.model_parallel_group)
1700
1701    def get_grad_norm_direct(self, gradients, params, norm_type=2):
1702        """Clips gradient norm of an iterable of parameters.
1703
1704        This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and
1705        added functionality to handle model parallel parameters. Note that
1706        the gradients are modified in place.
1707
1708        Arguments:
1709            parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
1710                single Tensor that will have gradients normalized
1711            max_norm (float or int): max norm of the gradients
1712            norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
1713                infinity norm.
1714
1715        Returns:
1716            Total norm of the parameters (viewed as a single vector).
1717        """
1718        norm_type = float(norm_type)
1719        all_norms = []
1720        if norm_type == inf:
1721            for g in gradients:
1722                all_norms.append(g.data.abs().max().float())
1723            total_norm = torch.stack(all_norms).max()
1724            dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=self.dp_process_group)
1725
1726            # Take max across all GPUs.
1727            self._model_parallel_all_reduce(tensor=total_norm, op=dist.ReduceOp.MAX)
1728        else:
1729            # if dist.get_rank() == 0:
1730            #    logger.info(f"Total Norm beginning {total_norm}")
1731            for g, p in zip(gradients, params):
1732                # Pipeline parallelism may replicate parameters. Avoid multi-counting.
1733                if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated:
1734                    continue
1735                if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0):
1736                    opguard_wrapped__all_norms_append(1688)(
1737                        all_norms.append, torch.norm(g.data.double().detach(), norm_type).to(get_accelerator().current_device_name()))
1738            if len(all_norms) > 0:
1739                total_norm = opguard_wrapped__total_norm_float(1691)(torch.stack(all_norms).square().sum().float)
1740            else:
1741                total_norm = torch.tensor(0.0, dtype=torch.float32).to(self.device)
1742            # Sum across all model parallel Device.
1743            dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=self.dp_process_group)
1744
1745            self._model_parallel_all_reduce(tensor=total_norm, op=dist.ReduceOp.SUM)
1746
1747            total_norm = opguard_wrapped__total_norm_pow(1699)(total_norm.pow, 1. / norm_type)
1748
1749        norm_is_inf = opguard_wrapped__norm_is_inf_isinf(1701)(total_norm.isinf)
1750        norm_is_nan = opguard_wrapped__norm_is_nan_isnan(1702)(total_norm.isnan)
1751        inf_or_nan = opguard_wrapped__inf_or_nan_logical_or(1703)(norm_is_nan.logical_or, norm_is_inf)
1752
1753        err = opguard_wrapped__err_tensor(1705)(torch.tensor, -1.0, device=self.device, dtype=torch.float)
1754        total_norm = inf_or_nan * err + opguard_wrapped__total_norm_logical_not(1706)(inf_or_nan.logical_not) * total_norm
1755        return total_norm
1756
1757    # creates a flat fused tensor from the tensor list starting at the first_offset
1758    # in the first tensor of the list. If there are not enough elements in the tensor
1759    # list then the flat tensor will be padded with zeros
1760    def get_flat_partition(self, tensor_list, first_offset, partition_size, dtype, device, return_tensor_list=False):
1761        flat_tensor_list = []
1762        current_size = 0
1763
1764        for i, tensor in enumerate(tensor_list):
1765            grad_accum = self.get_param_gradient_attribute(tensor)
1766            if grad_accum is None:
1767                grad_accum = torch.zeros_like(tensor, dtype=dtype)
1768
1769            tensor = grad_accum
1770            num_elements = tensor.numel()
1771            tensor_offset = 0
1772
1773            # we need to offset to get to the right element
1774            if i == 0 and first_offset > 0:
1775                tensor_offset = first_offset
1776                num_elements = num_elements - tensor_offset
1777
1778            # we dont need all elements of the tensor
1779            if num_elements > (partition_size - current_size):
1780                num_elements = partition_size - current_size
1781
1782            # we need a narrow view of the tensor based on the tensor offset and number of elements that
1783            # we need from this tensor
1784            if tensor_offset > 0 or num_elements < tensor.numel():
1785                flat_tensor_list.append(tensor.contiguous().view(-1).narrow(0, int(tensor_offset), int(num_elements)))
1786            else:
1787                flat_tensor_list.append(tensor)
1788
1789            current_size = current_size + num_elements
1790
1791        # this means its the last partition and does not align with the dp boundary. We need to pad before flattening
1792        if current_size < partition_size:
1793            flat_tensor_list.append(torch.zeros(int(partition_size - current_size), dtype=dtype, device=device))
1794
1795        if return_tensor_list:
1796            return flat_tensor_list
1797
1798        return self.flatten(flat_tensor_list)
1799
1800    def free_grad_in_param_list(self, param_list):
1801        for p in param_list:
1802            p.grad = None  # in step
1803            p.grad_accum = None
1804
1805    def reset_cpu_buffers(self):
1806        self.norm_for_param_grads = {}
1807        self.local_overflow = False
1808
1809    def set_lr(self, lr):
1810        """Set the learning rate."""
1811        for param_group in self.optimizer.param_groups:
1812            param_group["lr"] = lr
1813
1814    def get_lr(self):
1815        """Return the current learning rate."""
1816        return self.optimizer.param_groups[0]["lr"]
1817
1818    def override_loss_scale(self, loss_scale):
1819        if loss_scale != self.external_loss_scale:
1820            logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}')
1821        self.custom_loss_scaler = True
1822        self.external_loss_scale = loss_scale
1823
1824    def scaled_global_norm(self, norm_type=2):
1825        assert norm_type == 2, "only L2 norm supported"
1826        norm_groups = []
1827        for i, group in enumerate(self.bit16_groups):
1828            if self.cpu_offload:
1829                # complete complete_grad_norm_calculation_for_cpu_offload return python float, moving back to
1830                # torch.tensor as else statement returns tensor as well
1831                norm = torch.tensor(self.complete_grad_norm_calculation_for_cpu_offload(self.params_in_partition[i]),
1832                                    device=self.device)
1833                norm_groups.append(norm)
1834            else:
1835                norm_groups.append(self.get_grad_norm_direct(self.averaged_gradients[i], self.params_in_partition[i]))
1836
1837        if self.has_moe_layers:
1838            self._average_expert_grad_norms(norm_groups)
1839
1840        # calculating L2 norm
1841        return opguard_wrapped__stack_norm(1793)(torch.norm, torch.stack(norm_groups), p=norm_type)
1842
1843    def get_bit16_param_group(self, group_no):
1844        bit16_partitions = self.parallel_partitioned_bit16_groups[group_no]
1845        partition_id = dist.get_rank(group=self.real_dp_process_group[group_no])
1846        return [bit16_partitions[dist.get_rank(group=self.real_dp_process_group[group_no])]]
1847
1848    def _optimizer_step(self, group_no):
1849        original_param_groups = self.optimizer.param_groups
1850        self.optimizer.param_groups = [original_param_groups[group_no]]
1851        # Disabling this as the C++ side copy & synchronize is not working correctly
1852        #from deepspeed.ops.adam import DeepSpeedCPUAdam
1853        #if type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half:
1854        #    self.optimizer.step(fp16_param_groups=[self.get_bit16_param_group(group_no)])
1855        #else:
1856        #    self.optimizer.step()
1857        self.optimizer.step()
1858        self.optimizer.param_groups = original_param_groups
1859
1860        # We need to link optimizer state after the first step() call
1861        self._lazy_init_hp_params_optimizer_state()
1862
1863    def step(self, closure=None):
1864        """
1865        Not supporting closure.
1866        """
1867        self.micro_step_id = -1
1868
1869        see_memory_usage(f"In step before checking overflow")
1870
1871        # First compute norm for all group so we know if there is overflow
1872        if self.dtype == torch.float16:
1873            self.check_overflow()
1874
1875        prev_scale = self.loss_scale
1876        self._update_scale(self.overflow)
1877        if self.overflow:
1878            see_memory_usage('After overflow before clearing gradients')
1879            self.zero_grad(set_to_none=True)
1880            if self.cpu_offload:
1881                self.reset_cpu_buffers()
1882            else:
1883                self.averaged_gradients = {}
1884
1885            see_memory_usage('After overflow after clearing gradients')
1886
1887            for timer in OPTIMIZER_TIMERS:
1888                self.timers(timer).start()
1889                self.timers(timer).stop()
1890            return
1891
1892        # Step 1:- Calculate gradient norm using bit-16 grads
1893        see_memory_usage('Before norm calculation')
1894        scaled_global_grad_norm = self.scaled_global_norm()
1895        self._global_grad_norm = opguard_wrapped_scaled_global_grad_norm_truediv_prev_scale(1847)(operator.truediv, scaled_global_grad_norm, prev_scale)
1896        see_memory_usage('After norm before optimizer')
1897
1898        # Step 2:- run optimizer and upscaling simultaneously
1899        for i, group in enumerate(self.bit16_groups):
1900            self.timers(OPTIMIZER_GRADIENTS_TIMER).start()
1901            partition_id = dist.get_rank(group=self.real_dp_process_group[i])
1902            if self.cpu_offload:
1903                single_grad_partition = self.single_partition_of_fp32_groups[i].grad
1904                self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm)
1905
1906                self.timers(OPTIMIZER_GRADIENTS_TIMER).stop()
1907                self.timers(OPTIMIZER_STEP_TIMER).start()
1908                self._optimizer_step(i)
1909
1910                # Disabled, this is not currently working
1911                #from deepspeed.ops.adam import DeepSpeedCPUAdam
1912                #if not (type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half):
1913                #    bit16_partitions = self.parallel_partitioned_bit16_groups[i]
1914                #    fp32_partition = self.single_partition_of_fp32_groups[i]
1915                #    bit16_partitions[partition_id].data.copy_(fp32_partition.data)
1916                bit16_partitions = self.parallel_partitioned_bit16_groups[i]
1917                fp32_partition = self.single_partition_of_fp32_groups[i]
1918                opguard_wrapped__data_copy_(1870)(
1919                    bit16_partitions[partition_id].data.copy_, fp32_partition.to(get_accelerator().current_device_name()).data)
1920
1921                self.timers(OPTIMIZER_STEP_TIMER).stop()
1922            else:
1923                # free gradients for all the parameters that are not updated by this process(ZeRO stage2)
1924                self.free_grad_in_param_list(self.params_not_in_partition[i])
1925
1926                # create a flat gradients for parameters updated by this process
1927                # If we are last partition, ensure we have same size grads and partition size, if not pad with zero tensors
1928                if partition_id == dist.get_world_size(group=self.real_dp_process_group[i]) - 1:
1929                    single_grad_partition = opguard_wrapped__single_grad_partition_to(1881)(self.flatten_dense_tensors_aligned(
1930                        self.averaged_gradients[i],
1931                        int(self.partition_size[i])).to, self.single_partition_of_fp32_groups[i].dtype)
1932                else:
1933                    single_grad_partition = opguard_wrapped__single_grad_partition_to(1885)(
1934                        self.flatten(self.averaged_gradients[i]).to, self.single_partition_of_fp32_groups[i].dtype)
1935                assert single_grad_partition.numel() == self.partition_size[i], \
1936                    "averaged gradients have different number of elements that partition size {} {} {} {}".format(
1937                        single_grad_partition.numel(), self.partition_size[i], i, partition_id)
1938
1939                self.single_partition_of_fp32_groups[i].grad = single_grad_partition
1940                # release all the gradient since we have already created a necessary copy in dp_grad_partition(ZeRO stage2)
1941                self.free_grad_in_param_list(self.params_in_partition[i])
1942
1943                self.averaged_gradients[i] = None
1944
1945                self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm)
1946
1947                self.timers(OPTIMIZER_GRADIENTS_TIMER).stop()
1948
1949                # Step 3:- run the optimizer if no offloading
1950                self.timers(OPTIMIZER_STEP_TIMER).start()
1951                self._optimizer_step(i)
1952                # Step 4:- get rid of the fp32 gradients. Not needed anymore
1953                self.single_partition_of_fp32_groups[i].grad = None
1954                del single_grad_partition
1955                bit16_partitions = self.parallel_partitioned_bit16_groups[i]
1956                fp32_partition = self.single_partition_of_fp32_groups[i]
1957                opguard_wrapped__data_copy_(1909)(bit16_partitions[partition_id].data.copy_, fp32_partition.data)
1958                self.timers(OPTIMIZER_STEP_TIMER).stop()
1959
1960        see_memory_usage('After optimizer before all-gather')
1961        if self.cpu_offload:
1962            self.reset_cpu_buffers()
1963
1964        self.timers(OPTIMIZER_ALLGATHER_TIMER).start()
1965        # Gather the updated weights from everyone.
1966        # Then all partitions of the model parameters are updated and ready for next round forward.
1967        all_gather_dp_groups(groups_flat=self.bit16_groups_flat,
1968                             partitioned_param_groups=self.parallel_partitioned_bit16_groups,
1969                             dp_process_group=self.real_dp_process_group,
1970                             start_alignment_factor=self.nccl_start_alignment_factor,
1971                             allgather_bucket_size=self.allgather_bucket_size)
1972        self.timers(OPTIMIZER_ALLGATHER_TIMER).stop()
1973
1974        # TODO: we probably don't need this? just to be safe
1975        for i in range(len(self.bit16_groups)):
1976            self._update_model_bit16_weights(i)
1977
1978        self.timers.log(OPTIMIZER_TIMERS)
1979        see_memory_usage('After zero_optimizer step')
1980
1981        return
1982
1983    @torch.no_grad()
1984    def update_lp_params(self):
1985        for i, (bit16_partitions, fp32_partition) in enumerate(
1986                zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)):
1987            partition_id = dist.get_rank(group=self.real_dp_process_group[i])
1988            opguard_wrapped__data_copy_(1940)(bit16_partitions[partition_id].data.copy_, fp32_partition.data)
1989            # print_rank_0(f'update_lp_params {i=} {partition_id=}', force=True)
1990            # if i == 0:
1991            #     print_rank_0(f'{fp32_partition[:10]=}', force=True)
1992        all_gather_dp_groups(groups_flat=self.bit16_groups_flat,
1993                             partitioned_param_groups=self.parallel_partitioned_bit16_groups,
1994                             dp_process_group=self.real_dp_process_group,
1995                             start_alignment_factor=self.nccl_start_alignment_factor,
1996                             allgather_bucket_size=self.allgather_bucket_size)
1997
1998    def _average_expert_grad_norms(self, norm_groups):
1999        for i, norm in enumerate(norm_groups):
2000            if self.is_moe_param_group[i]:
2001                scaled_norm_tensor = norm * 1.0 / dist.get_world_size(group=self.real_dp_process_group[i])
2002                if self.device == 'cpu':
2003                    scaled_norm_tensor = scaled_norm_tensor.to(get_accelerator().current_device_name())
2004                dist.all_reduce(scaled_norm_tensor, group=self.real_dp_process_group[i])
2005                norm_groups[i] = scaled_norm_tensor.to(self.device)
2006
2007    def unscale_and_clip_grads(self, grad_groups_flat, total_norm):
2008        # compute combined scale factor for this group
2009        combined_scale = self.loss_scale
2010        if self.clip_grad > 0.:
2011            # norm is in fact norm*scale
2012            clip = opguard_wrapped_total_norm_truediv_clip_grad(1964)(operator.truediv, ((total_norm / self.loss_scale) + 1e-6), self.clip_grad)
2013            if clip > 1:
2014                combined_scale = clip * self.loss_scale
2015
2016        for grad in grad_groups_flat:
2017            if isinstance(grad, list):
2018                sub_partitions = grad
2019                for g in sub_partitions:
2020                    g.data.mul_(1. / combined_scale)
2021            else:
2022                opguard_wrapped__data_mul_(1974)(grad.data.mul_, 1. / combined_scale)
2023
2024    def _check_overflow(self, partition_gradients=True):
2025        self.overflow = self.has_overflow(partition_gradients)
2026
2027    # `params` is a list / generator of torch.Variable
2028    def has_overflow_serial(self, params):
2029        invalid_grad_count = opguard_wrapped__invalid_grad_count_zeros(1981)(torch.zeros, [1], dtype=torch.float, device=get_accelerator().current_device_name())
2030        for p in params:
2031            if p.grad is not None:
2032                invalid_grad_count += opguard_wrapped__invalid_grad_count__has_inf_or_nan(1984)(self._has_inf_or_nan, p.grad)
2033        return opguard_wrapped__invalid_grad_count_bool(1985)(invalid_grad_count.bool)
2034
2035    def has_overflow_partitioned_grads_serial(self):
2036        invalid_grad_count = opguard_wrapped__invalid_grad_count_zeros(1988)(torch.zeros, [1], dtype=torch.float, device=get_accelerator().current_device_name())
2037        for i in range(len(self.bit16_groups)):
2038            for j, grad in enumerate(self.averaged_gradients[i]):
2039                if grad is not None:
2040                    invalid_grad_count += opguard_wrapped__invalid_grad_count__has_inf_or_nan(1992)(self._has_inf_or_nan, grad)
2041        return opguard_wrapped__invalid_grad_count_bool(1993)(invalid_grad_count.bool)
2042
2043    def has_overflow(self, partition_gradients=True):
2044        if partition_gradients:
2045            overflow = self.local_overflow if self.cpu_offload else self.has_overflow_partitioned_grads_serial()
2046            overflow_gpu = get_accelerator().ByteTensor([overflow]) if self.cpu_offload else opguard_wrapped__overflow_gpu_to(1998)(
2047                overflow.byte().to, get_accelerator().current_device_name())
2048            '''This will capture overflow across all data parallel and expert parallel process
2049            Since expert parallel process are a subset of data parallel process'''
2050            dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.dp_process_group)
2051
2052        else:
2053            params = []
2054            for group in self.bit16_groups:
2055                for param in group:
2056                    params.append(param)
2057            overflow_gpu = opguard_wrapped__overflow_gpu_to(2009)(self.has_overflow_serial(params).byte().to, get_accelerator().current_device_name())
2058
2059        # Since each model parallel GPU carries only part of the model,
2060        # make sure overflow flag is synced across all the model parallel GPUs
2061        self._model_parallel_all_reduce(tensor=overflow_gpu, op=dist.ReduceOp.MAX)
2062
2063        overflow = opguard_wrapped__overflow_item(2015)(overflow_gpu[0].item)
2064        return bool(overflow)
2065
2066    # `x` is a torch.Tensor
2067    @staticmethod
2068    def _has_inf_or_nan(x, j=None):
2069        float_x = opguard_wrapped__float_x_float(2021)(x.float)
2070        nan = opguard_wrapped__nan_isnan(2022)(float_x.isnan)
2071        inf = opguard_wrapped__inf_isinf(2023)(float_x.isinf)
2072        inf_or_nan = opguard_wrapped__inf_or_nan_logical_or(2024)(nan.logical_or, inf)
2073        return opguard_wrapped__float_max(2025)(inf_or_nan.float().max)
2074
2075    def backward(self, loss, retain_graph=False):
2076        """
2077        :attr:`backward` performs the following steps:
2078
2079        1. fp32_loss = loss.float()
2080        2. scaled_loss = fp32_loss*loss_scale
2081        3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves
2082        """
2083        self.micro_step_id += 1
2084
2085        if self.contiguous_gradients:
2086            self.ipg_buffer = []
2087            buf_0 = opguard_wrapped__buf_0_empty(2039)(torch.empty, int(self.reduce_bucket_size),
2088                                dtype=self.dtype,
2089                                device=get_accelerator().current_device_name())
2090            self.ipg_buffer.append(buf_0)
2091
2092            # Use double buffers to avoid data access conflict when overlap_comm is enabled.
2093            if self.overlap_comm:
2094                buf_1 = opguard_wrapped__buf_1_empty(2046)(torch.empty, int(self.reduce_bucket_size),
2095                                    dtype=self.dtype,
2096                                    device=get_accelerator().current_device_name())
2097                self.ipg_buffer.append(buf_1)
2098            self.ipg_index = 0
2099
2100        if self.custom_loss_scaler:
2101            scaled_loss = self.external_loss_scale * loss
2102            scaled_loss.backward()
2103        else:
2104            self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
2105
2106        # Only for Stage 1, Mode 2
2107        if self.use_grad_accum_attribute:
2108            self.fill_grad_accum_attribute()
2109
2110    def check_overflow(self, partition_gradients=True):
2111        self._check_overflow(partition_gradients)
2112
2113    def _update_scale(self, has_overflow=False):
2114        self.loss_scaler.update_scale(has_overflow)
2115
2116    # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state"
2117    def _get_state(self):
2118        return self.optimizer.state
2119
2120    def _set_state(self, value):
2121        self.optimizer.state = value
2122
2123    state = property(_get_state, _set_state)
2124
2125    # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups"
2126    # (for example, to adjust the learning rate)
2127    def _get_param_groups(self):
2128        return self.optimizer.param_groups
2129
2130    def _set_param_groups(self, value):
2131        self.optimizer.param_groups = value
2132
2133    param_groups = property(_get_param_groups, _set_param_groups)
2134
2135    # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale"
2136    def _get_loss_scale(self):
2137        if self.custom_loss_scaler:
2138            return self.external_loss_scale
2139        else:
2140            return self.loss_scaler.cur_scale
2141
2142    def _set_loss_scale(self, value):
2143        self.loss_scaler.cur_scale = value
2144
2145    loss_scale = property(_get_loss_scale, _set_loss_scale)
2146    cur_scale = property(_get_loss_scale, _set_loss_scale)
2147
2148    # Return group tensor after removing paddings that are added for alignment to DP world size.
2149    # This method works on the assumption that each group contains a single flattened tensor.
2150    def _get_groups_without_padding(self, groups_with_padding):
2151        groups_without_padding = []
2152        for i, group in enumerate(groups_with_padding):
2153            lean_length = group.numel() - self.groups_padding[i]
2154            groups_without_padding.append(group[:lean_length])
2155
2156        return groups_without_padding
2157
2158    # Return optimizer state after removing paddings that are added for alignment.
2159    def _get_state_without_padding(self, state_with_padding, padding):
2160        lean_state = {}
2161        for key, value in state_with_padding.items():
2162            if torch.is_tensor(value):
2163                lean_length = value.numel() - padding
2164                lean_state[key] = value[:lean_length]
2165            else:
2166                lean_state[key] = value
2167
2168        return lean_state
2169
2170    # Return base optimizer states.
2171    # This method assumes that each param group contains a single flattened tensor.
2172    def _get_base_optimizer_state(self):
2173        optimizer_groups_state = []
2174        for i, group in enumerate(self.optimizer.param_groups):
2175            p = group['params'][0]
2176            lean_optimizer_state = self._get_state_without_padding(self.optimizer.state[p], self.groups_padding[i])
2177            optimizer_groups_state.append(lean_optimizer_state)
2178
2179        return optimizer_groups_state
2180
2181    def state_dict(self):
2182        """
2183        Returns a dict containing the current state of this :class:`FP16_Optimizer` instance.
2184        This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict
2185        of the contained Pytorch optimizer.
2186        Example::
2187            checkpoint = {}
2188            checkpoint['model'] = model.state_dict()
2189            checkpoint['optimizer'] = optimizer.state_dict()
2190            torch.save(checkpoint, "saved.pth")
2191        """
2192        state_dict = {}
2193        state_dict[LOSS_SCALER] = self.loss_scaler
2194        state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale
2195        state_dict['overflow'] = self.overflow
2196        state_dict[CLIP_GRAD] = self.clip_grad
2197
2198        if self.elastic_checkpoint:
2199            state_dict[BASE_OPTIMIZER_STATE] = self._get_base_optimizer_state()
2200
2201            if "step" in self.optimizer.param_groups[0]:
2202                # Assuming "step" is the only item that changes through training iterations
2203                assert all(group["step"] == self.optimizer.param_groups[0]["step"]
2204                           for group in self.optimizer.param_groups), "All param groups must have the same step value"
2205                state_dict[BASE_OPTIMIZER_STATE_STEP] = self.optimizer.param_groups[0]["step"]
2206        else:
2207            state_dict[BASE_OPTIMIZER_STATE] = self.optimizer.state_dict()
2208
2209        # Remove paddings for DP alignment to enable loading for other alignment values
2210        fp32_groups_without_padding = self._get_groups_without_padding(self.single_partition_of_fp32_groups)
2211        state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = fp32_groups_without_padding
2212
2213        state_dict[
2214            ZERO_STAGE] = ZeroStageEnum.gradients if self.partition_gradients else ZeroStageEnum.optimizer_states
2215        state_dict[GROUP_PADDINGS] = self.groups_padding
2216        state_dict[PARTITION_COUNT] = self.partition_count
2217
2218        state_dict[DS_VERSION] = version
2219        state_dict[PARAM_SLICE_MAPPINGS] = self._param_slice_mappings
2220
2221        return state_dict
2222
2223    # Restore base optimizer fp32 weights from elastic checkpoint by:
2224    # 1) Merging fp32 weights from checkpoints of all partitions
2225    # 2) Extracting fp32 weights for current partition from merged weights
2226    # 3) Using extracted weights to update base optimizer weights directly.
2227    def _restore_from_elastic_fp32_weights(self, all_state_dict):
2228        merged_single_partition_of_fp32_groups = []
2229
2230        for i in range(len(self.single_partition_of_fp32_groups)):
2231            partition_id = dist.get_rank(group=self.real_dp_process_group[i])
2232            merged_partitions = [sd[SINGLE_PARTITION_OF_FP32_GROUPS][i] for sd in all_state_dict]
2233            if self.is_moe_group(self.optimizer.param_groups[i]):
2234                ranks = self.get_ep_ranks(group_name=self.optimizer.param_groups[i]['name'])
2235                merged_partitions = [merged_partitions[i] for i in ranks]
2236            flat_merged_partitions = self.flatten_dense_tensors_aligned(
2237                merged_partitions,
2238                self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[i]))
2239            dp_partitions = self.get_data_parallel_partitions(flat_merged_partitions, i)
2240            merged_single_partition_of_fp32_groups.append(dp_partitions[partition_id])
2241
2242        for current, saved in zip(self.single_partition_of_fp32_groups, merged_single_partition_of_fp32_groups):
2243            current.data.copy_(saved.data)
2244
2245    # Restore base optimizer fp32 weights from ZeRO fp16 or bfloat16 weights
2246    def _restore_from_bit16_weights(self):
2247        for group_id, (bit16_partitions, fp32_partition) in enumerate(
2248                zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)):
2249            partition_id = dist.get_rank(group=self.real_dp_process_group[group_id])
2250            fp32_partition.data.copy_(bit16_partitions[partition_id].data)
2251
2252    # Refresh the fp32 master params from the fp16 or bfloat16 copies.
2253    def refresh_fp32_params(self):
2254        self._restore_from_bit16_weights()
2255
2256    # Extract optimizer state for current partition from merged states of all partitions
2257    def _partition_base_optimizer_state(self, state_key, all_partition_states, group_id):
2258        partition_id = dist.get_rank(group=self.real_dp_process_group[group_id])
2259        alignment = self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[group_id])
2260        if torch.is_tensor(all_partition_states[0]):
2261            flat_merged_partitions = self.flatten_dense_tensors_aligned(all_partition_states, alignment)
2262            dp_partitions = self.get_data_parallel_partitions(flat_merged_partitions, group_id)
2263            return dp_partitions[partition_id]
2264        else:
2265            # Assume non-tensor states are not partitioned and equal across ranks, so return first one
2266            return all_partition_states[0]
2267
2268    def _restore_step_from_elastic_checkpoint(self, all_state_dict):
2269        assert BASE_OPTIMIZER_STATE_STEP in all_state_dict[0]
2270        assert all(sd[BASE_OPTIMIZER_STATE_STEP] == all_state_dict[0][BASE_OPTIMIZER_STATE_STEP]
2271                   for sd in all_state_dict), "State dicts of all partitions must have the same step value"
2272        return all_state_dict[0][BASE_OPTIMIZER_STATE_STEP]
2273
2274    def _restore_base_optimizer_state(self, base_optimizer_group_states, base_optimizer_state_step, group_paddings):
2275        if type(base_optimizer_group_states) == dict:
2276            base_optimizer_group_states = base_optimizer_group_states['state']
2277
2278        saved_keys = base_optimizer_group_states[0].keys()
2279
2280        for i, group in enumerate(self.optimizer.param_groups):
2281            p = group['params'][0]
2282            padding = 0 if group_paddings is None else group_paddings[i]
2283            for key in saved_keys:
2284                saved = base_optimizer_group_states[i][key]
2285
2286                if torch.is_tensor(saved):
2287                    if key in self.optimizer.state[p]:
2288                        dst_tensor = self.optimizer.state[p][key]
2289                        src_tensor = _get_padded_tensor(saved, dst_tensor.numel())
2290                        self.optimizer.state[p][key].data.copy_(src_tensor.data)
2291                    else:
2292                        self.optimizer.state[p][key] = _pad_tensor_by_size(
2293                            saved, padding, torch.float32,
2294                            torch.device('cpu') if self.cpu_offload else self.device)
2295                else:
2296                    self.optimizer.state[p][key] = saved
2297
2298        for param_group in self.optimizer.param_groups:
2299            param_group['step'] = base_optimizer_state_step
2300
2301    def get_ep_ranks(self, rank=0, group_name=None):
2302        from deepspeed.utils import groups
2303        expert_parallel_size_ = groups._get_expert_parallel_world_size(group_name)
2304        world_size = groups._get_data_parallel_world_size()
2305        rank = groups._get_expert_parallel_rank(group_name)
2306        ranks = range(rank, world_size, expert_parallel_size_)
2307        return list(ranks)
2308
2309    # Restore base optimizer state from elastic checkpoint by
2310    # 1) Merging optimizer state from checkpoints of all partitions
2311    # 2) Extracting optimizer state for current partition from the merged state
2312    # 3) Using the extracted value to directly update the base optimizer.
2313    def _restore_elastic_base_optimizer_state(self, all_state_dict):
2314        base_optimizer_group_states = []
2315        for i in range(len(self.optimizer.param_groups)):
2316            partition_states = {}
2317            all_partition_group_states = [sd[BASE_OPTIMIZER_STATE][i] for sd in all_state_dict]
2318
2319            if self.is_moe_group(self.optimizer.param_groups[i]):
2320                ranks = self.get_ep_ranks(group_name=self.optimizer.param_groups[i]['name'])
2321                all_partition_group_states = [all_partition_group_states[i] for i in ranks]
2322
2323            for key in all_partition_group_states[0].keys():
2324                all_partition_states = [all_states[key] for all_states in all_partition_group_states]
2325                partition_states[key] = self._partition_base_optimizer_state(key, all_partition_states, i)
2326            base_optimizer_group_states.append(partition_states)
2327
2328        self._restore_base_optimizer_state(base_optimizer_group_states,
2329                                           self._restore_step_from_elastic_checkpoint(all_state_dict), None)
2330
2331    def load_state_dict(self,
2332                        state_dict_list,
2333                        load_optimizer_states=True,
2334                        load_from_fp32_weights=False,
2335                        checkpoint_folder=None,
2336                        load_serial=None):
2337        if checkpoint_folder:
2338            self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights)
2339        else:
2340            self._load_legacy_checkpoint(state_dict_list, load_optimizer_states, load_from_fp32_weights)
2341
2342    def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights):
2343        self.load_hp_checkpoint_state_from_checkpoint_dir("bit16_groups", checkpoint_folder)
2344
2345    @property
2346    def param_groups(self):
2347        """Forward the wrapped optimizer's parameters."""
2348        return self.optimizer.param_groups
2349
2350    def _load_global_state(self, sd):
2351        self.loss_scaler = sd.get(LOSS_SCALER, self.loss_scaler)
2352        self.dynamic_loss_scale = sd.get('dynamic_loss_scale', self.dynamic_loss_scale)
2353        self.overflow = sd.get('overflow', self.overflow)
2354        self.clip_grad = sd.get(CLIP_GRAD, self.clip_grad)
2355
2356        ckpt_version = sd.get(DS_VERSION, False)
2357        assert ckpt_version, f"Empty ds_version in checkpoint, not clear how to proceed"
2358        ckpt_version = pkg_version.parse(ckpt_version)
2359
2360        # zero stage 1 mode
2361        if not self.partition_gradients:
2362            required_version = pkg_version.parse("0.3.17")
2363            error_str = f"ZeRO stage 1 changed in {required_version} and is not backwards compatible " \
2364                "with older stage 1 checkpoints. If you'd like to load an old ZeRO-1 checkpoint " \
2365                "please use an older version of DeepSpeed (<= 0.5.8) and set 'legacy_stage1': true in your zero config json."
2366            assert required_version <= ckpt_version, f"Old version: {ckpt_version} {error_str}"
2367
2368    def _load_legacy_checkpoint(self, state_dict_list, load_optimizer_states=True, load_from_fp32_weights=False):
2369        r"""Loading ZeRO checkpoint
2370
2371        Arguments:
2372            state_dict_list: List of all saved ZeRO checkpoints, one for each saved partition.
2373                Note that the number of saved partitions may differ from number of loading partitions to support
2374                changing GPU count, specifically DP world size, between saving and loading checkpoints.
2375            load_optimizer_states: Boolean indicating whether or not to load base optimizer states
2376            load_from_fp32_weights: Boolean indicating whether to initialize fp32 master weights from fp32
2377            copies in checkpoints (no precision loss) or from model's fp16 copies (with precision loss).
2378        """
2379        """
2380        Loads a state_dict created by an earlier call to state_dict().
2381        If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``,
2382        whose parameters in turn came from ``model``, it is expected that the user
2383        will call ``model.load_state_dict()`` before
2384        ``fp16_optimizer_instance.load_state_dict()`` is called.
2385        Example::
2386            model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half()
2387            optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
2388            optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
2389            ...
2390            checkpoint = torch.load("saved.pth")
2391            model.load_state_dict(checkpoint['model'])
2392            optimizer.load_state_dict(checkpoint['optimizer'])
2393        """
2394
2395        # I think it should actually be ok to reload the optimizer before the model.
2396        dp_rank = dist.get_rank(group=self.dp_process_group)
2397        current_rank_sd = state_dict_list[dp_rank]
2398        self._load_global_state(current_rank_sd)
2399
2400        ckpt_is_rigid = isinstance(current_rank_sd[BASE_OPTIMIZER_STATE], dict)
2401
2402        # padding is always at the last rank/partition
2403        # if DP=1024 and param-group elems=16 -> padding will be 1024-16 across all but one rank
2404        # scenario-1 (shrink): saving w. 4 gpus -> loading w. 2 gpus
2405        # scenario-2 (expand): saving w. 2 gpus -> loading w. 4 gpus
2406        # if load_optimizer_states:
2407        #     if new_dp_size:
2408        #         self.strip_padding()
2409        #         self.add_padding_w_new_dp_size()
2410        #     self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE])
2411
2412        if load_optimizer_states:
2413            if ckpt_is_rigid:
2414                # loading rigid ckpt into either rigid or elastic exec
2415                self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE])
2416            else:
2417                if self.elastic_checkpoint:
2418                    # loading elastic into elastic exec
2419                    self._restore_elastic_base_optimizer_state(state_dict_list)
2420                else:
2421                    # loading an elastic checkpoint into rigid exec
2422                    self._restore_base_optimizer_state(current_rank_sd[BASE_OPTIMIZER_STATE],
2423                                                       current_rank_sd[BASE_OPTIMIZER_STATE_STEP],
2424                                                       current_rank_sd[GROUP_PADDINGS])
2425
2426        # At this point, the optimizer's references to the model's fp32 parameters are up to date.
2427        # The optimizer's hyperparameters and internal buffers are also up to date.
2428        # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still
2429        # out of date.  There are two options.
2430        # 1:  Refresh the master params from the model's fp16 params.
2431        # This requires less storage but incurs precision loss.
2432        # 2:  Save and restore the fp32 master copies separately.
2433        # We choose option 1 if changing DP degree and option 2 otherwise.
2434        #
2435        # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device
2436        # of their associated parameters, because it's possible those buffers might not exist yet in
2437        # the current optimizer instance.  In our case, as long as the current FP16_Optimizer has been
2438        # constructed in the same way as the one whose state_dict we are loading, the same master params
2439        # are guaranteed to exist, so we can just copy_() from the saved master params.
2440
2441        if load_from_fp32_weights:
2442            # option 2 from above
2443            if self.elastic_checkpoint and not ckpt_is_rigid:
2444                self._restore_from_elastic_fp32_weights(state_dict_list)
2445            else:
2446                # For non-elastic checkpoint, simply copying from saved weights of current rank is sufficient.
2447                for current, saved in zip(self.single_partition_of_fp32_groups,
2448                                          current_rank_sd[SINGLE_PARTITION_OF_FP32_GROUPS]):
2449                    src_tensor = _get_padded_tensor(saved, current.numel())
2450                    current.data.copy_(src_tensor.data)
2451        else:
2452            # option 1 from above
2453            self._restore_from_bit16_weights()
2454
2455        if load_optimizer_states:
2456            self._link_all_hp_params()
2457
2458
2459def _handle_overflow(cpu_sum, x, i):
2460    import math
2461    rank = dist.get_rank()
2462    if rank == 0:
2463        t_i = -1
2464        for v_i, v in enumerate(x.data.contiguous().view(-1)):
2465            if not math.isfinite(float(v)):
2466                t_i = v_i
2467                break
2468        logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}")
2469
2470
2471def estimate_zero2_model_states_mem_needs(total_params,
2472                                          num_gpus_per_node=1,
2473                                          num_nodes=1,
2474                                          cpu_offload=True,
2475                                          additional_buffer_factor=1.5):
2476
2477    total_gpus = num_nodes * num_gpus_per_node
2478
2479    if cpu_offload:
2480        gpu_mem = 2 * total_params
2481        cpu_mem = total_params * max(4 * total_gpus, 16) * additional_buffer_factor
2482    else:
2483        gpu_mem = 4 * total_params + int(16 * total_params / total_gpus)
2484        cpu_mem = total_params * 4 * num_gpus_per_node * additional_buffer_factor
2485
2486    return int(cpu_mem), int(gpu_mem)
2487
2488
2489def model_to_params(model):
2490    # shared params calculated only once
2491    total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())
2492    return total_params
2493
2494
2495def estimate_zero2_model_states_mem_needs_all_live(model,
2496                                                   num_gpus_per_node=1,
2497                                                   num_nodes=1,
2498                                                   additional_buffer_factor=1.5):
2499    """
2500    Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients
2501    for a given ``model`` and hardware setup.
2502
2503    If you have an actual model object, use this function and everything will be derived
2504    automatically.
2505
2506    If it's a hypothetical model, use ``estimate_zero2_model_states_mem_needs_all_cold`` where you have to pass
2507    the ``total_params`` explicitly.
2508
2509    Args:
2510        - ``model``: ``nn.Module`` object
2511        - ``num_gpus_per_node``: how many gpus per node (defaults to 1)
2512        - ``num_nodes``: how many nodes (defaults to 1),
2513        - ``additional_buffer_factor``: estimation factor (defaults to 1.5):
2514
2515    """
2516
2517    total_params = model_to_params(model)
2518
2519    estimate_zero2_model_states_mem_needs_all_cold(total_params=total_params,
2520                                                   num_gpus_per_node=num_gpus_per_node,
2521                                                   num_nodes=num_nodes,
2522                                                   additional_buffer_factor=additional_buffer_factor)
2523
2524
2525def estimate_zero2_model_states_mem_needs_all_cold(total_params,
2526                                                   num_gpus_per_node=1,
2527                                                   num_nodes=1,
2528                                                   additional_buffer_factor=1.5):
2529    """
2530    Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients
2531    for a given ``model`` and hardware setup.
2532
2533    If it's a hypothetical model, use this function where you have to pass
2534    the ``total_params`` and ``largest_layer_params`` explicitly.
2535
2536    If you have an actual model object, use ``estimate_zero2_model_states_mem_needs_all_live`` and everything
2537    will be derived automatically.
2538
2539    Args:
2540        - ``total_params``: total  model params
2541        - ``num_gpus_per_node``: how many gpus per node (defaults to 1)
2542        - ``num_nodes``: how many nodes (defaults to 1),
2543        - ``additional_buffer_factor``: estimation factor (defaults to 1.5):
2544
2545    """
2546
2547    def format_options(cpu_offload):
2548        enabled = []
2549        device = f'{OffloadDeviceEnum.cpu:4}' if cpu_offload else "none"
2550        enabled.append(f"offload_optimizer={device}")
2551        return ", ".join(enabled)
2552
2553    nodes_str = "nodes" if num_nodes > 1 else "node"
2554    gpus_str = "GPUs" if num_gpus_per_node > 1 else "GPU"
2555    print("Estimated memory needed for params, optim states and gradients for a:\n"
2556          f"HW: Setup with {num_nodes} {nodes_str}, {num_gpus_per_node} {gpus_str} per node.\n"
2557          f"SW: Model with {int(total_params/1e6)}M total params.")
2558    print("  per CPU  |  per GPU |   Options")
2559    for cpu_offload in [True, False]:
2560        cpu_mem, gpu_mem = estimate_zero2_model_states_mem_needs(total_params=total_params,
2561                                                                 num_gpus_per_node=num_gpus_per_node,
2562                                                                 num_nodes=num_nodes,
2563                                                                 cpu_offload=cpu_offload,
2564                                                                 additional_buffer_factor=additional_buffer_factor)
2565
2566        options_str = format_options(cpu_offload=cpu_offload)
2567        print(f" {cpu_mem/2**30:7.2f}GB | {gpu_mem/2**30:6.2f}GB | {options_str}")