| """Load full policy weights before or after DeepSpeed partitions parameters.""" | |
| def load_policy_state_dict(model, state_dict, *, strict=True): | |
| """Every distributed rank must call this with the same state dictionary. | |
| ZeRO-3 initialization replaces parameter storage with empty shards. Gather | |
| those parameters while copying weights, then broadcast and repartition on | |
| context exit. Loading on every rank also initializes replicated buffers. | |
| Ordinary models do not need DeepSpeed installed. | |
| """ | |
| partitioned = [parameter for parameter in model.parameters() if hasattr(parameter, "ds_id")] | |
| if not partitioned: | |
| return model.load_state_dict(state_dict, strict=strict) | |
| from deepspeed import zero | |
| with zero.GatheredParameters(partitioned, modifier_rank=0): | |
| return model.load_state_dict(state_dict, strict=strict) | |