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)
| 5225 | return frozen_param_fragments |
| 5226 | |
| 5227 | def _get_zero_param_shapes(self): |
| 5228 | """Returns a dict of name to shape mapping, only for the flattened fp32 weights saved by the |
| 5229 | optimizer. the names are exactly as in state_dict. The order is absolutely important, since |
| 5230 | the saved data is just flattened data with no identifiers and requires reconstruction in the |
| 5231 | same order it was saved. |
| 5232 | We can't rely on self.module.named_parameters() to get the saved tensors, as some params |
| 5233 | will be missing and others unsaved and then it'd be impossible to reconstruct state_dict |
| 5234 | from the flattened weights. |
| 5235 | optimizer.bit16_groups seems to be the easiest to use as it's in all zeroX versions. |
| 5236 | """ |
| 5237 | param_group_shapes = [] |
| 5238 | cnt = 0 |
| 5239 | numel = 0 |
| 5240 | |
| 5241 | # zero2 started using a round_robin_bit16_groups which is a shuffled version of bit16_groups - |
| 5242 | # if we don't use it, we get parameters ordered incorrectly |
| 5243 | if hasattr(self.optimizer, "round_robin_bit16_groups"): |
| 5244 | bit16_groups = self.optimizer.round_robin_bit16_groups |
| 5245 | elif self.bfloat16_enabled() and hasattr(self.optimizer, "bf16_groups"): |
| 5246 | bit16_groups = self.optimizer.bf16_groups |
| 5247 | else: |
| 5248 | bit16_groups = self.optimizer.bit16_groups if self.zero_optimization_stage( |
| 5249 | ) == 2 else self.optimizer.fp16_groups |
| 5250 | |
| 5251 | for bit16_group in bit16_groups: |
| 5252 | param_shapes = OrderedDict() |
| 5253 | for param in bit16_group: |
| 5254 | cnt += 1 |
| 5255 | numel += param.ds_numel if hasattr(param, "ds_numel") else param.numel() |
| 5256 | shape = param.ds_shape if hasattr(param, "ds_shape") else param.shape |
| 5257 | if param not in self.param_names: |
| 5258 | raise ValueError("failed to find optimizer param in named params") |
| 5259 | name = self.param_names[param] |
| 5260 | param_shapes[name] = shape |
| 5261 | |
| 5262 | # uncomment to debug zero_to_fp32.py problems |
| 5263 | # if self.global_rank == 0: print(f"saving param {name} {shape} (numel={shape.numel()})") |
| 5264 | param_group_shapes.append(param_shapes) |
| 5265 | # if self.global_rank == 0: print(f"Total saved {numel} numels in {cnt} params") |
| 5266 | |
| 5267 | return param_group_shapes |
| 5268 | |
| 5269 | def _get_shared_params(self): |
| 5270 | """ |