| """Write a model-only copy of an FSDP2 DCP checkpoint and verify it bit for bit. |
| |
| Each __R_0.distcp shard holds rank R's slice of both the model and the Adam |
| state, so the optimizer cannot be removed by deleting files. This reads every |
| "model.*" tensor (fp32) on CPU without a process group, saves them under the |
| same "model" prefix the trainer and eval/hrm_loader.py use, reloads the new |
| copy and compares every tensor with torch.equal. |
| |
| python tools/export_model_only_dcp.py <src fsdp2_epoch_N> <dst dir> |
| """ |
| import sys |
| import time |
|
|
| import torch |
| import torch.distributed.checkpoint as dcp |
| from torch.distributed.checkpoint.metadata import TensorStorageMetadata |
|
|
| src, dst = sys.argv[1], sys.argv[2] |
| t0 = time.time() |
| md = dcp.FileSystemReader(src).read_metadata().state_dict_metadata |
| model_md = {k: v for k, v in md.items() if k.startswith("model.")} |
| non_tensor = [k for k, v in model_md.items() if not isinstance(v, TensorStorageMetadata)] |
| if non_tensor: |
| sys.exit(f"FATAL: non-tensor model entries, refusing: {non_tensor[:5]}") |
| state = {k[len("model."):]: torch.empty(v.size, dtype=v.properties.dtype) for k, v in model_md.items()} |
| dcp.load({"model": state}, checkpoint_id=src, no_dist=True) |
| dcp.save({"model": state}, checkpoint_id=dst, no_dist=True) |
|
|
| check = {k: torch.empty_like(v) for k, v in state.items()} |
| dcp.load({"model": check}, checkpoint_id=dst, no_dist=True) |
| bad = [k for k in state if not torch.equal(state[k], check[k])] |
| if bad: |
| sys.exit(f"FATAL: {len(bad)} tensors differ after round trip: {bad[:5]}") |
| n = sum(v.numel() for v in state.values()) |
| print(f"OK {len(state)} tensors, {n:,} params, {n * 4 / 2**30:.2f} GiB fp32, bit-exact, {time.time() - t0:.0f}s") |
|
|
