merge.py

import torch
import torch.distributed.tensor
from safetensors.torch import save_file
import os
from collections import OrderedDict
import gc

def merge_fsdp_to_safetensors(rank0_path, rank1_path, output_path, target_dtype=None):
    """
    FSDP로 분할된 두 개의 .pt 파일을 하나의 .safetensors 파일로 병합
    
    Args:
        rank0_path (str): rank 0 .pt 파일 경로
        rank1_path (str): rank 1 .pt 파일 경로  
        output_path (str): 출력할 .safetensors 파일 경로
        target_dtype (torch.dtype, optional): 타겟 dtype (예: torch.float16, torch.bfloat16)
    """
    print("Loading rank 0 checkpoint...")
    rank0_dict = torch.load(rank0_path, map_location='cpu', weights_only=False)
    
    print("Loading rank 1 checkpoint...")
    rank1_dict = torch.load(rank1_path, map_location='cpu', weights_only=False)
    
    # DTensor를 일반 텐서로 변환하는 함수
    def convert_dtensor_to_tensor(state_dict):
        converted_dict = OrderedDict()
        dtype_info = {}
        for key, value in state_dict.items():
            if hasattr(value, '_local_tensor'):
                # DTensor인 경우 로컬 텐서 추출
                tensor = value._local_tensor
                converted_dict[key] = tensor
                dtype_info[key] = tensor.dtype
                print(f"Converted DTensor to tensor: {key} (dtype: {tensor.dtype})")
            elif isinstance(value, torch.Tensor):
                converted_dict[key] = value
                dtype_info[key] = value.dtype
            else:
                # 다른 타입은 그대로 유지
                converted_dict[key] = value
                dtype_info[key] = type(value).__name__
        return converted_dict, dtype_info
    
    print("Converting DTensors to regular tensors...")
    rank0_dict, rank0_dtypes = convert_dtensor_to_tensor(rank0_dict)
    rank1_dict, rank1_dtypes = convert_dtensor_to_tensor(rank1_dict)
    
    # dtype 정보 출력
    print("\n📋 Original dtype information:")
    all_dtypes_r0 = set(dtype_info for dtype_info in rank0_dtypes.values() if isinstance(dtype_info, torch.dtype))
    all_dtypes_r1 = set(dtype_info for dtype_info in rank1_dtypes.values() if isinstance(dtype_info, torch.dtype))
    all_dtypes = all_dtypes_r0 | all_dtypes_r1
    
    print(f"   Rank 0 dtypes found: {all_dtypes_r0}")
    print(f"   Rank 1 dtypes found: {all_dtypes_r1}")
    print(f"   All dtypes: {all_dtypes}")
    
    if target_dtype:
        print(f"   Target dtype specified: {target_dtype}")
    else:
        print("   No target dtype specified - keeping original dtypes")
    
    # 병합된 상태 사전
    merged_state_dict = OrderedDict()
    
    # rank 0의 모든 키들을 먼저 처리
    all_keys = set(rank0_dict.keys()) | set(rank1_dict.keys())
    
    print(f"Total unique keys found: {len(all_keys)}")
    
    for key in sorted(all_keys):
        rank0_tensor = rank0_dict.get(key)
        rank1_tensor = rank1_dict.get(key)
        
        if rank0_tensor is not None and rank1_tensor is not None:
            # 두 rank에 모두 존재하는 경우 - 차원 확인 후 연결
            print(f"Merging key: {key}")
            print(f"  Rank 0 shape: {rank0_tensor.shape}, dtype: {rank0_tensor.dtype}")
            print(f"  Rank 1 shape: {rank1_tensor.shape}, dtype: {rank1_tensor.dtype}")
            
            # dtype 변환 (필요한 경우)
            if target_dtype and rank0_tensor.dtype != target_dtype:
                rank0_tensor = rank0_tensor.to(target_dtype)
                print(f"  Converted rank 0 to {target_dtype}")
            if target_dtype and rank1_tensor.dtype != target_dtype:
                rank1_tensor = rank1_tensor.to(target_dtype)
                print(f"  Converted rank 1 to {target_dtype}")
            
            # 첫 번째 차원으로 연결 (일반적인 FSDP 샤딩 방식)
            merged_tensor = torch.cat([rank0_tensor, rank1_tensor], dim=0)
            merged_state_dict[key] = merged_tensor
            print(f"  Merged shape: {merged_tensor.shape}, dtype: {merged_tensor.dtype}")
            
        elif rank0_tensor is not None:
            # rank 0에만 존재
            tensor = rank0_tensor
            if target_dtype and isinstance(tensor, torch.Tensor) and tensor.dtype != target_dtype:
                tensor = tensor.to(target_dtype)
                print(f"Converting {key} from rank 0: {rank0_tensor.dtype} -> {target_dtype}")
            print(f"Adding from rank 0: {key} (shape: {tensor.shape if isinstance(tensor, torch.Tensor) else 'N/A'}, dtype: {tensor.dtype if isinstance(tensor, torch.Tensor) else type(tensor).__name__})")
            merged_state_dict[key] = tensor
            
        elif rank1_tensor is not None:
            # rank 1에만 존재
            tensor = rank1_tensor
            if target_dtype and isinstance(tensor, torch.Tensor) and tensor.dtype != target_dtype:
                tensor = tensor.to(target_dtype)
                print(f"Converting {key} from rank 1: {rank1_tensor.dtype} -> {target_dtype}")
            print(f"Adding from rank 1: {key} (shape: {tensor.shape if isinstance(tensor, torch.Tensor) else 'N/A'}, dtype: {tensor.dtype if isinstance(tensor, torch.Tensor) else type(tensor).__name__})")
            merged_state_dict[key] = tensor
    
    print(f"\nTotal merged parameters: {len(merged_state_dict)}")
    
    # 메모리 정리
    del rank0_dict, rank1_dict
    gc.collect()
    
    # safetensors로 저장
    print(f"Saving merged model to {output_path}...")
    
    # 최종 dtype 정보 출력
    final_dtypes = {}
    for key, tensor in merged_state_dict.items():
        if isinstance(tensor, torch.Tensor):
            final_dtypes[tensor.dtype] = final_dtypes.get(tensor.dtype, 0) + 1
    
    print(f"📋 Final merged model dtype distribution:")
    for dtype, count in final_dtypes.items():
        print(f"   {dtype}: {count} tensors")
    
    save_file(merged_state_dict, output_path)
    print("✅ Successfully saved merged model!")
    
    return merged_state_dict

def merge_with_custom_concatenation(rank0_path, rank1_path, output_path, concat_rules=None):
    """
    사용자 정의 연결 규칙으로 병합
    
    Args:
        concat_rules (dict): 키별 연결 차원 지정 {'key_pattern': dim}
    """
    if concat_rules is None:
        # 기본 규칙
        concat_rules = {
            'weight': 0,  # 가중치는 첫 번째 차원으로 연결
            'bias': 0,    # 편향도 첫 번째 차원으로 연결
        }
    
    print("Loading checkpoints...")
    rank0_dict = torch.load(rank0_path, map_location='cpu', weights_only=False)
    rank1_dict = torch.load(rank1_path, map_location='cpu', weights_only=False)
    
    merged_state_dict = OrderedDict()
    all_keys = set(rank0_dict.keys()) | set(rank1_dict.keys())
    
    for key in sorted(all_keys):
        rank0_tensor = rank0_dict.get(key)
        rank1_tensor = rank1_dict.get(key)
        
        if rank0_tensor is not None and rank1_tensor is not None:
            # 연결 차원 결정
            concat_dim = 0  # 기본값
            for pattern, dim in concat_rules.items():
                if pattern in key:
                    concat_dim = dim
                    break
            
            print(f"Merging {key} along dimension {concat_dim}")
            merged_tensor = torch.cat([rank0_tensor, rank1_tensor], dim=concat_dim)
            merged_state_dict[key] = merged_tensor
            
        elif rank0_tensor is not None:
            merged_state_dict[key] = rank0_tensor
        elif rank1_tensor is not None:
            merged_state_dict[key] = rank1_tensor
    
    # 정리 및 저장
    del rank0_dict, rank1_dict
    gc.collect()
    
    print(f"Saving to {output_path}...")
    save_file(merged_state_dict, output_path)
    print("✅ Merge completed!")

def comprehensive_verification(rank0_path, rank1_path, merged_path):
    """병합이 올바르게 되었는지 종합적으로 검증"""
    import torch.distributed.tensor
    from safetensors import safe_open
    
    print("🔍 Starting comprehensive verification...")
    
    # 1. 원본 파일들 로드
    print("\n📁 Loading original files...")
    rank0_dict = torch.load(rank0_path, map_location='cpu', weights_only=False)
    rank1_dict = torch.load(rank1_path, map_location='cpu', weights_only=False)
    
    # DTensor를 일반 텐서로 변환
    def convert_dtensor_to_tensor(state_dict):
        converted_dict = {}
        for key, value in state_dict.items():
            if hasattr(value, '_local_tensor'):
                converted_dict[key] = value._local_tensor
            elif isinstance(value, torch.Tensor):
                converted_dict[key] = value
            else:
                converted_dict[key] = value
        return converted_dict
    
    rank0_dict = convert_dtensor_to_tensor(rank0_dict)
    rank1_dict = convert_dtensor_to_tensor(rank1_dict)
    
    # 2. 원본 파일들 분석
    rank0_keys = set(rank0_dict.keys())
    rank1_keys = set(rank1_dict.keys())
    all_original_keys = rank0_keys | rank1_keys
    shared_keys = rank0_keys & rank1_keys
    rank0_only = rank0_keys - rank1_keys
    rank1_only = rank1_keys - rank0_keys
    
    print(f"📊 Original files analysis:")
    print(f"   Rank 0 keys: {len(rank0_keys)}")
    print(f"   Rank 1 keys: {len(rank1_keys)}")
    print(f"   Shared keys: {len(shared_keys)}")
    print(f"   Rank 0 only: {len(rank0_only)}")
    print(f"   Rank 1 only: {len(rank1_only)}")
    print(f"   Total unique keys: {len(all_original_keys)}")
    
    # 3. 원본 파라미터 수 계산
    original_params = 0
    original_shapes = {}
    
    for key in all_original_keys:
        if key in shared_keys:
            # 공유 키는 두 텐서를 연결한 크기로 계산
            r0_tensor = rank0_dict[key]
            r1_tensor = rank1_dict[key]
            combined_shape = list(r0_tensor.shape)
            combined_shape[0] += r1_tensor.shape[0]  # 첫 번째 차원으로 연결 가정
            original_shapes[key] = tuple(combined_shape)
            original_params += r0_tensor.numel() + r1_tensor.numel()
        elif key in rank0_only:
            original_shapes[key] = rank0_dict[key].shape
            original_params += rank0_dict[key].numel()
        elif key in rank1_only:
            original_shapes[key] = rank1_dict[key].shape
            original_params += rank1_dict[key].numel()
    
    print(f"   Original total parameters: {original_params:,}")
    
    # 4. 병합된 파일 분석
    print(f"\n📁 Loading merged file: {merged_path}")
    merged_params = 0
    merged_keys = set()
    merged_shapes = {}
    
    with safe_open(merged_path, framework="pt", device="cpu") as f:
        merged_keys = set(f.keys())
        for key in f.keys():
            tensor = f.get_tensor(key)
            merged_shapes[key] = tensor.shape
            merged_params += tensor.numel()
    
    print(f"📊 Merged file analysis:")
    print(f"   Merged keys: {len(merged_keys)}")
    print(f"   Merged parameters: {merged_params:,}")
    
    # 5. 비교 및 검증
    print(f"\n✅ Verification Results:")
    
    # 키 개수 비교
    keys_match = len(merged_keys) == len(all_original_keys)
    print(f"   Keys count match: {keys_match} ({len(merged_keys)} vs {len(all_original_keys)})")
    
    # 파라미터 수 비교
    params_match = merged_params == original_params
    print(f"   Parameter count match: {params_match} ({merged_params:,} vs {original_params:,})")
    
    # 키 이름 비교
    missing_keys = all_original_keys - merged_keys
    extra_keys = merged_keys - all_original_keys
    
    if missing_keys:
        print(f"   ❌ Missing keys: {missing_keys}")
    
    if extra_keys:
        print(f"   ❌ Extra keys: {extra_keys}")
    
    # 개별 텐서 크기 비교
    shape_mismatches = []
    for key in merged_keys & all_original_keys:
        if merged_shapes[key] != original_shapes[key]:
            shape_mismatches.append((key, merged_shapes[key], original_shapes[key]))
    
    if shape_mismatches:
        print(f"   ❌ Shape mismatches:")
        for key, merged_shape, original_shape in shape_mismatches[:5]:  # 처음 5개만 표시
            print(f"      {key}: {merged_shape} vs {original_shape}")
        if len(shape_mismatches) > 5:
            print(f"      ... and {len(shape_mismatches) - 5} more")
    
    # 6. 세부 분석 (선택적)
    print(f"\n📋 Detailed Analysis:")
    print(f"   Shared keys that should be concatenated:")
    for key in sorted(list(shared_keys))[:10]:  # 처음 10개만 표시
        r0_shape = rank0_dict[key].shape
        r1_shape = rank1_dict[key].shape
        expected_shape = list(r0_shape)
        expected_shape[0] += r1_shape[0]
        actual_shape = merged_shapes.get(key, "MISSING")
        status = "✅" if tuple(expected_shape) == actual_shape else "❌"
        print(f"      {status} {key}: {r0_shape} + {r1_shape} -> {actual_shape}")
    
    if len(shared_keys) > 10:
        print(f"      ... and {len(shared_keys) - 10} more shared keys")
    
    # 7. 최종 결과
    overall_success = keys_match and params_match and not missing_keys and not extra_keys and not shape_mismatches
    
    print(f"\n{'='*50}")
    if overall_success:
        print("🎉 MERGE VERIFICATION SUCCESSFUL!")
        print("   All parameters have been correctly merged.")
    else:
        print("⚠️  MERGE VERIFICATION FOUND ISSUES!")
        print("   Please review the mismatches above.")
    print(f"{'='*50}")
    
    # 정리
    del rank0_dict, rank1_dict
    gc.collect()
    
    return overall_success

# 사용 예시
if __name__ == "__main__":
    # 파일 경로 설정
    rank0_file = "model_rank_0.pt"  # 실제 파일명으로 변경
    rank1_file = "model_rank_1.pt"  # 실제 파일명으로 변경
    output_file = "merged_model.safetensors"
    
    # dtype 옵션 설정
    target_dtype = torch.bfloat16  # bf16으로 변환

    # 기본 병합
    print("Starting merge process...")
    merged_dict = merge_fsdp_to_safetensors(rank0_file, rank1_file, output_file, target_dtype)

    # 종합적인 검증
    print("\nStarting comprehensive verification...")
    verification_passed = comprehensive_verification(rank0_file, rank1_file, output_file)

    if verification_passed:
        print(f"\n🎉 Successfully merged and verified FSDP model to {output_file}")
    else:
        print(f"\n⚠️  Merge completed but verification found issues. Please review the output above.")

    # 추가: 간단한 로드 테스트
    print(f"\n🔍 Testing if merged model can be loaded...")
    try:
        from safetensors import safe_open
        with safe_open(output_file, framework="pt", device="cpu") as f:
            sample_keys = list(f.keys())[:3]
            for key in sample_keys:
                tensor = f.get_tensor(key)
                print(f"   ✅ Successfully loaded {key}: {tensor.shape}, dtype: {tensor.dtype}")
        print("   ✅ Merged model loads correctly!")
    except Exception as e:
        print(f"   ❌ Error loading merged model: {e}")
        
    print(f"\n💡 Tip: To change dtype, modify 'target_dtype' in the script:")
    print(f"   - torch.float16 for fp16 (smaller file, less precision)")
    print(f"   - torch.bfloat16 for bf16 (good balance)")
    print(f"   - torch.float32 for fp32 (larger file, best precision)")
    print(f"   - None to keep original dtypes")