[tab] tools/merge_fsdp_checkpoint.py (python)
import os
import torch
import torch.distributed.checkpoint as dcp
from transformers import AutoConfig, AutoModelForCausalLM
import safetensors.torch
import json
from pathlib import Path

# Mocking ModelState as used in miles
class ModelState:
    def __init__(self, model):
        self.model = model
    def state_dict(self):
        # Return a nested dict to match the checkpoint structure
        return {"model": self.model.state_dict()}
    def load_state_dict(self, state_dict):
        self.model.load_state_dict(state_dict["model"])

def merge_checkpoint(checkpoint_dir, output_dir, hf_source_dir):
    checkpoint_dir = Path(checkpoint_dir)
    output_dir = Path(output_dir)
    hf_source_dir = Path(hf_source_dir)
    
    model_dir = checkpoint_dir / "model"
    
    print(f"Loading config from {hf_source_dir}")
    config = AutoConfig.from_pretrained(hf_source_dir, trust_remote_code=True)
    
    print("Initialising empty model (on CPU)...")
    with torch.device("meta"):
        model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
    
    # Materialise on CPU (or zero-init)
    model = model.to_empty(device="cpu")
    
    # The checkpoint has a "model_state" key at the top level because of how dcp.save was called in miles
    # state_dict = {"model_state": model_state}
    # And model_state.state_dict() returns {"model": model.state_dict()}
    # So the full path to a param is model_state.model.model.<param_name>
    
    # We create a nested structure to match dcp's expectations
    # Based on peek_checkpoint_metadata.py: 
    # model_state.model.model.embed_tokens.weight
    
    # This suggests dcp.load expects a state_dict that HAS those keys or a plan that maps them.
    # dcp.load(state_dict=state_dict, ...)
    
    class SimpleStateful:
        def __init__(self, model):
            self.model = model
        def state_dict(self):
            # This needs to be exactly what was saved.
            # In miles/backends/fsdp_utils/checkpoint.py:
            # model_state = ModelState(actor.model) -> returns {"model": model_state_dict}
            # state_dict = {"model_state": model_state}
            # In ModelState.state_dict(): model_state_dict, _ = get_state_dict(self.model)
            # FSDP get_state_dict typically returns the UNWRAPPED keys if configured correctly.
            # But the metadata showed "model_state.model.model.embed_tokens.weight"
            # This means it's Double wrapped: model_state -> model -> model
            
            # Let's try to match the nesting
            return {"model": {"model": self.model.state_dict()}}
        def load_state_dict(self, sd):
            self.model.load_state_dict(sd["model"]["model"])

    state_dict = {"model_state": SimpleStateful(model)}
    
    print(f"Loading sharded weights from {model_dir}...")
    # no_dist=True allows loading sharded checkpoints on a single process
    dcp.load(state_dict=state_dict, checkpoint_id=str(model_dir), no_dist=True)
    
    print(f"Saving merged model to {output_dir}...")
    output_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(output_dir, safe_serialization=True)
    
    # Copy other files
    import shutil
    for f in hf_source_dir.glob("*"):
        if f.is_file() and not f.name.endswith(".safetensors") and not f.name.endswith(".bin") and not f.name.endswith(".pt"):
            shutil.copy(f, output_dir / f.name)
            
    print("Merge complete!")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint-dir", type=str, required=True)
    parser.add_argument("--output-dir", type=str, required=True)
    parser.add_argument("--hf-source-dir", type=str, required=True)
    args = parser.parse_args()
    
    merge_checkpoint(args.checkpoint_dir, args.output_dir, args.hf_source_dir)