"""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)