import argparse import os import torch from transformers import AutoModel def convert(model_name, output_path, tensor_parallel_size, use_te): """Convert InternViT HF checkpoint to mcore.""" hf_model = AutoModel.from_pretrained( model_name, trust_remote_code=True ) hf_state_dict = hf_model.state_dict() new_state_dicts = [{"model": dict()} for _ in range(tensor_parallel_size)] hidden_size = 3200 num_heads = 25 dim = 128 order = torch.ones(3 * hidden_size).long() for j in range(num_heads): for i in range(dim): order[i + dim*3*j] = j*dim+i order[dim + i + dim*3*j] = j*dim+i+num_heads*dim order[dim*2 + i + dim*3*j] = j*dim+i+num_heads*dim*2 for name, tensor in hf_state_dict.items(): # Map parameter names to ones used in megatron. new_name = "" new_tensor = tensor # This is used for chunking some tensors to target tensor parallel size. chunk_dim = None if "embeddings.class_embedding" in name: new_name = "class_token" elif "embeddings.patch_embedding.weight" in name: new_name = "conv1.weight" elif "embeddings.patch_embedding.bias" in name: new_name = "conv1.bias" elif "embeddings.position_embedding" in name: new_name = "position_embeddings.weight" new_tensor = new_tensor.squeeze(0) elif "encoder.layers" in name: layer_idx = name.split(".")[2] base = f"decoder.layers.{layer_idx}" head_dim = 128 if tensor_parallel_size == 1: num_padded_heads = 25 elif tensor_parallel_size == 8: # Note: 25 is not divisible by 8 and we don't currently support uneven heads split with tensor parallelism. # So we pad with dummy all-zero heads. Please use a nice even number of attention heads in your model. num_padded_heads = 32 else: raise NotImplementedError("invalid tensor parallel size value:", tensor_parallel_size) if "ls1" in name: new_name = f"{base}.ls1" elif "ls2" in name: new_name = f"{base}.ls2" elif "attn.qkv.weight" in name: new_name = f"{base}.self_attention.linear_qkv.weight" num_tensors = 3 padded_dim = head_dim * num_padded_heads * num_tensors padded_tensor = torch.zeros((padded_dim, new_tensor.shape[-1]), dtype=new_tensor.dtype, device=new_tensor.device) padded_tensor[:new_tensor.shape[0], :] = new_tensor[order] new_tensor = padded_tensor chunk_dim = 0 elif "attn.q_norm.weight" in name: new_name = f"{base}.self_attention.q_layernorm.weight" num_tensors = 1 padded_dim = head_dim * num_padded_heads * num_tensors padded_tensor = torch.zeros(padded_dim, dtype=new_tensor.dtype, device=new_tensor.device) padded_tensor[:new_tensor.shape[0]] = new_tensor new_tensor = padded_tensor chunk_dim = 0 elif "attn.k_norm.weight" in name: new_name = f"{base}.self_attention.k_layernorm.weight" num_tensors = 1 padded_dim = head_dim * num_padded_heads * num_tensors padded_tensor = torch.zeros(padded_dim, dtype=new_tensor.dtype, device=new_tensor.device) padded_tensor[:new_tensor.shape[0]] = new_tensor new_tensor = padded_tensor chunk_dim = 0 elif "attn.proj.weight" in name: new_name = f"{base}.self_attention.linear_proj.weight" num_tensors = 1 padded_dim = head_dim * num_padded_heads * num_tensors padded_tensor = torch.zeros((new_tensor.shape[0], padded_dim), dtype=new_tensor.dtype, device=new_tensor.device) padded_tensor[:, :new_tensor.shape[-1]] = new_tensor new_tensor = padded_tensor chunk_dim = 1 elif "attn.proj.bias" in name: new_name = f"{base}.self_attention.linear_proj.bias" elif "mlp.fc1.weight" in name: new_name = f"{base}.mlp.linear_fc1.weight" chunk_dim = 0 elif "mlp.fc1.bias" in name: new_name = f"{base}.mlp.linear_fc1.bias" chunk_dim = 0 elif "mlp.fc2.weight" in name: new_name = f"{base}.mlp.linear_fc2.weight" chunk_dim = 1 elif "mlp.fc2.bias" in name: new_name = f"{base}.mlp.linear_fc2.bias" elif "norm1" in name: new_name = f"{base}.input_layernorm.weight" elif "norm2" in name: new_name = f"{base}.pre_mlp_layernorm.weight" else: raise RuntimeError("unexpected transformer layer name", name) else: raise RuntimeError("unexpected layer name", name) assert new_name != "", f"unexpected layer name {name}" # TE sets _extra_state (for FP8 purposes), so set an empty one here for compatibility. extra_state_layers = ("linear_qkv", "linear_proj", "linear_fc1", "linear_fc2") is_extra_state_layer = any([l in new_name for l in extra_state_layers]) if use_te and is_extra_state_layer: layer = new_name.split(".")[-2] if layer in extra_state_layers: extra_state_name = ( new_name[: new_name.rfind(".") + 1] + "_extra_state" ) # Replace the weight name. for i in range(tensor_parallel_size): new_state_dicts[i]["model"][extra_state_name] = None if chunk_dim is None: new_tensors = [new_tensor for _ in range(tensor_parallel_size)] else: new_tensors = torch.chunk(new_tensor, tensor_parallel_size, dim=chunk_dim) for i in range(tensor_parallel_size): new_state_dicts[i]["model"][new_name] = new_tensors[i].clone() for i in range(tensor_parallel_size): output_dir_tp = os.path.join(output_path, f"iter_0000001/mp_rank_0{i}") os.makedirs(output_dir_tp, exist_ok=True) output_path_tp = os.path.join(output_dir_tp, "model_optim_rng.pt") torch.save(new_state_dicts[i], output_path_tp) print("saved file", output_path_tp) print("done") if __name__ == "__main__": parser = argparse.ArgumentParser(description="InternVIT HuggingFace to Mcore converter") parser.add_argument("--model-name", type=str, default="OpenGVLab/InternViT-6B-448px-V1-5", help="Model name in HuggingFace") parser.add_argument("--output-dir", type=str, required=True, help="Output directory for the mcore model.") parser.add_argument("--use-te", action="store_true", default=True) parser.add_argument("--tensor-parallel-size", type=int, required=True) args = parser.parse_args() convert(args.model_name, args.output_dir, args.tensor_parallel_size, args.use_te)