Returns a dict of name to shape mapping, only for the flattened fp32 weights saved by the optimizer. the names are exactly as in state_dict. The order is absolutely important, since the saved data is just flattened data with no identifiers and requires reconstruction in the s
(self)
| 5100 | return frozen_param_fragments |
| 5101 | |
| 5102 | def _get_zero_param_shapes(self): |
| 5103 | """Returns a dict of name to shape mapping, only for the flattened fp32 weights saved by the |
| 5104 | optimizer. the names are exactly as in state_dict. The order is absolutely important, since |
| 5105 | the saved data is just flattened data with no identifiers and requires reconstruction in the |
| 5106 | same order it was saved. |
| 5107 | We can't rely on self.module.named_parameters() to get the saved tensors, as some params |
| 5108 | will be missing and others unsaved and then it'd be impossible to reconstruct state_dict |
| 5109 | from the flattened weights. |
| 5110 | optimizer.bit16_groups seems to be the easiest to use as it's in all zeroX versions. |
| 5111 | """ |
| 5112 | param_group_shapes = [] |
| 5113 | cnt = 0 |
| 5114 | numel = 0 |
| 5115 | |
| 5116 | # zero2 started using a round_robin_bit16_groups which is a shuffled version of bit16_groups - |
| 5117 | # if we don't use it, we get parameters ordered incorrectly |
| 5118 | if hasattr(self.optimizer, "round_robin_bit16_groups"): |
| 5119 | bit16_groups = self.optimizer.round_robin_bit16_groups |
| 5120 | elif self.bfloat16_enabled() and hasattr(self.optimizer, "bf16_groups"): |
| 5121 | bit16_groups = self.optimizer.bf16_groups |
| 5122 | else: |
| 5123 | bit16_groups = self.optimizer.bit16_groups if self.zero_optimization_stage( |
| 5124 | ) == 2 else self.optimizer.fp16_groups |
| 5125 | |
| 5126 | for bit16_group in bit16_groups: |
| 5127 | param_shapes = OrderedDict() |
| 5128 | for param in bit16_group: |
| 5129 | cnt += 1 |
| 5130 | numel += param.ds_numel if hasattr(param, "ds_numel") else param.numel() |
| 5131 | shape = param.ds_shape if hasattr(param, "ds_shape") else param.shape |
| 5132 | if param not in self.param_names: |
| 5133 | raise ValueError("failed to find optimizer param in named params") |
| 5134 | name = self.param_names[param] |
| 5135 | param_shapes[name] = shape |
| 5136 | |
| 5137 | # uncomment to debug zero_to_fp32.py problems |
| 5138 | # if self.global_rank == 0: print(f"saving param {name} {shape} (numel={shape.numel()})") |
| 5139 | param_group_shapes.append(param_shapes) |
| 5140 | # if self.global_rank == 0: print(f"Total saved {numel} numels in {cnt} params") |
| 5141 | |
| 5142 | return param_group_shapes |
| 5143 | |
| 5144 | def _get_shared_params(self): |
| 5145 | """ |