#!/usr/bin/env python3
from __future__ import annotations

from pathlib import Path

import torch
import torch.distributed as dist
import torch.distributed.checkpoint as dcp


CKPT = Path(
    "runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/"
    "checkpoints/iter_0107500"
)


def init_dist() -> None:
    if not dist.is_initialized():
        dist.init_process_group("gloo", init_method="tcp://127.0.0.1:29631", rank=0, world_size=1)


def main() -> None:
    init_dist()
    reader = dcp.FileSystemReader(str(CKPT))
    metadata = reader.read_metadata()
    keys = list(metadata.state_dict_metadata.keys())
    print("num_keys", len(keys))
    for key in keys:
        item = metadata.state_dict_metadata[key]
        if hasattr(item, "size"):
            print("TENSOR", key, tuple(item.size), item.properties.dtype)

    target = "embedding.word_embeddings.weight"
    meta = metadata.state_dict_metadata[target]
    shape = tuple(meta.size)
    tensor = torch.empty(shape, dtype=torch.bfloat16)
    state = {target: tensor}
    dcp.load(state, checkpoint_id=str(CKPT))
    print("loaded", target, state[target].shape, state[target].dtype, state[target].float().mean().item())


if __name__ == "__main__":
    main()
