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