tasks/model_loading_example.py

#!/usr/bin/env python3
"""
示例：手动重现 lm-evaluation-harness 的模型加载流程

这个脚本展示了 lm-evaluation-harness 如何从 model_args 中读取配置并构建模型。
你可以运行这个脚本来理解整个流程。
"""

import sys
import os

# 添加 lm-evaluation-harness 到路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../lm-evaluation-harness'))

import torch
import transformers
from lm_eval import utils
from lm_eval.api.registry import get_model


def step1_parse_args_string():
    """步骤1: 解析 model_args 字符串"""
    print("=" * 80)
    print("步骤1: 解析 model_args 字符串")
    print("=" * 80)
    
    # 模拟命令行参数: --model_args "pretrained=EleutherAI/gpt-j-6B,dtype=float,revision=main"
    model_args_string = "pretrained=EleutherAI/gpt-j-6B,dtype=float,revision=main"
    
    print(f"输入字符串: {model_args_string}")
    
    # 使用 lm-eval 的解析函数
    parsed_args = utils.simple_parse_args_string(model_args_string)
    
    print(f"解析后的字典: {parsed_args}")
    print()
    
    return parsed_args


def step2_get_model_class():
    """步骤2: 获取模型类"""
    print("=" * 80)
    print("步骤2: 获取模型类")
    print("=" * 80)
    
    # 模拟命令行参数: --model hf
    model_name = "hf"
    
    print(f"模型名称: {model_name}")
    
    # 从注册表获取模型类
    model_class = get_model(model_name)
    
    print(f"获取到的模型类: {model_class}")
    print(f"类名: {model_class.__name__}")
    print()
    
    return model_class


def step3_load_config(pretrained, revision="main", trust_remote_code=False):
    """步骤3: 加载模型配置"""
    print("=" * 80)
    print("步骤3: 加载模型配置")
    print("=" * 80)
    
    print(f"模型路径/名称: {pretrained}")
    print(f"Revision: {revision}")
    print(f"Trust remote code: {trust_remote_code}")
    
    # 使用 transformers.AutoConfig.from_pretrained 加载配置
    # 这是 lm-eval 在 _get_config() 中做的事情
    config = transformers.AutoConfig.from_pretrained(
        pretrained,
        revision=revision,
        trust_remote_code=trust_remote_code,
    )
    
    print(f"\n配置类型: {type(config)}")
    print(f"模型类型: {config.model_type}")
    print(f"词汇表大小: {config.vocab_size}")
    print(f"隐藏层大小: {getattr(config, 'n_embd', getattr(config, 'hidden_size', 'N/A'))}")
    print(f"层数: {getattr(config, 'n_layer', getattr(config, 'num_hidden_layers', 'N/A'))}")
    print(f"注意力头数: {getattr(config, 'n_head', getattr(config, 'num_attention_heads', 'N/A'))}")
    
    # 显示配置文件的路径
    if hasattr(config, '_name_or_path'):
        print(f"\n配置来源: {config._name_or_path}")
    
    print()
    return config


def step4_load_tokenizer(pretrained, revision="main", trust_remote_code=False):
    """步骤4: 加载 tokenizer"""
    print("=" * 80)
    print("步骤4: 加载 Tokenizer")
    print("=" * 80)
    
    print(f"模型路径/名称: {pretrained}")
    
    # 使用 transformers.AutoTokenizer.from_pretrained 加载 tokenizer
    # 这是 lm-eval 在 _create_tokenizer() 中做的事情
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        pretrained,
        revision=revision,
        trust_remote_code=trust_remote_code,
        use_fast=True,
    )
    
    print(f"Tokenizer 类型: {type(tokenizer)}")
    print(f"词汇表大小: {tokenizer.vocab_size}")
    print(f"特殊 token:")
    print(f"  - BOS: {tokenizer.bos_token} (ID: {tokenizer.bos_token_id})")
    print(f"  - EOS: {tokenizer.eos_token} (ID: {tokenizer.eos_token_id})")
    print(f"  - PAD: {tokenizer.pad_token} (ID: {tokenizer.pad_token_id})")
    
    print()
    return tokenizer


def step5_load_model(pretrained, config, dtype="float", revision="main", trust_remote_code=False):
    """步骤5: 加载模型"""
    print("=" * 80)
    print("步骤5: 加载模型")
    print("=" * 80)
    
    print(f"模型路径/名称: {pretrained}")
    print(f"数据类型: {dtype}")
    
    # 转换 dtype 字符串为 torch.dtype
    if dtype == "float":
        torch_dtype = torch.float32
    elif dtype == "float16":
        torch_dtype = torch.float16
    elif dtype == "bfloat16":
        torch_dtype = torch.bfloat16
    else:
        torch_dtype = torch.float32
    
    print(f"Torch dtype: {torch_dtype}")
    
    # 根据配置确定模型类
    # 这是 lm-eval 在 _get_backend() 和 _create_model() 中做的事情
    model_type = config.model_type
    
    # 选择对应的模型类
    if model_type in ["gpt2", "gptj", "gpt_neox", "llama", "mistral", "mixtral"]:
        model_class = transformers.AutoModelForCausalLM
        print(f"使用模型类: AutoModelForCausalLM (因果语言模型)")
    elif model_type in ["t5", "bart"]:
        model_class = transformers.AutoModelForSeq2SeqLM
        print(f"使用模型类: AutoModelForSeq2SeqLM (序列到序列模型)")
    else:
        model_class = transformers.AutoModelForCausalLM
        print(f"使用默认模型类: AutoModelForCausalLM")
    
    print(f"\n开始加载模型（这可能需要一些时间）...")
    
    # 使用 from_pretrained 加载模型
    # 这是 lm-eval 在 _create_model() 中做的事情
    try:
        model = model_class.from_pretrained(
            pretrained,
            revision=revision,
            torch_dtype=torch_dtype,
            trust_remote_code=trust_remote_code,
            device_map="auto",  # 自动分配设备
        )
        
        print(f"模型类型: {type(model)}")
        print(f"模型参数数量: {sum(p.numel() for p in model.parameters()):,}")
        print(f"可训练参数数量: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
        
        # 显示模型设备
        if hasattr(model, 'device'):
            print(f"模型设备: {model.device}")
        elif hasattr(model, 'hf_device_map'):
            print(f"模型设备映射: {model.hf_device_map}")
        
    except Exception as e:
        print(f"加载模型时出错: {e}")
        print("注意: 这可能是由于内存不足或模型文件不存在")
        model = None
    
    print()
    return model


def step6_create_hflm_instance():
    """步骤6: 使用 lm-eval 创建完整的 HFLM 实例"""
    print("=" * 80)
    print("步骤6: 使用 lm-eval 创建完整的 HFLM 实例")
    print("=" * 80)
    
    # 模拟完整的命令行参数
    model_name = "hf"
    model_args_string = "pretrained=EleutherAI/gpt-j-6B,dtype=float"
    
    print(f"模型名称: {model_name}")
    print(f"模型参数: {model_args_string}")
    
    # 获取模型类
    model_class = get_model(model_name)
    
    # 使用 create_from_arg_string 创建实例（这是 lm-eval 实际使用的方法）
    print("\n使用 lm-eval 的 create_from_arg_string 创建模型实例...")
    
    try:
        # 注意: 实际创建模型需要时间和资源
        # 这里只是展示流程，不实际加载大模型
        print("（跳过实际模型加载，因为这需要大量内存和时间）")
        print("实际代码:")
        print(f"  lm = {model_class.__name__}.create_from_arg_string(")
        print(f"      '{model_args_string}',")
        print(f"      {{'batch_size': 8, 'device': 'cuda'}}")
        print(f"  )")
        
    except Exception as e:
        print(f"创建实例时出错: {e}")
    
    print()


def main():
    """主函数：演示完整的模型加载流程"""
    print("\n" + "=" * 80)
    print("lm-evaluation-harness 模型加载流程演示")
    print("=" * 80 + "\n")
    
    # 步骤1: 解析参数
    parsed_args = step1_parse_args_string()
    
    # 步骤2: 获取模型类
    model_class = step2_get_model_class()
    
    # 步骤3: 加载配置
    pretrained = parsed_args.get("pretrained", "EleutherAI/gpt-j-6B")
    revision = parsed_args.get("revision", "main")
    trust_remote_code = parsed_args.get("trust_remote_code", False)
    
    print("\n注意: 以下步骤需要网络连接来下载模型文件")
    print("如果模型已缓存，将使用缓存版本\n")
    
    try:
        config = step3_load_config(pretrained, revision, trust_remote_code)
        
        # 步骤4: 加载 tokenizer
        tokenizer = step4_load_tokenizer(pretrained, revision, trust_remote_code)
        
        # 步骤5: 加载模型（可选，需要大量内存）
        dtype = parsed_args.get("dtype", "float")
        
        print("\n是否加载完整模型？这需要大量内存和时间。")
        print("（在实际使用中，lm-eval 会自动执行这一步）")
        load_full_model = False  # 设置为 True 来实际加载模型
        
        if load_full_model:
            model = step5_load_model(pretrained, config, dtype, revision, trust_remote_code)
        else:
            print("跳过完整模型加载（设置 load_full_model=True 来实际加载）")
            print()
        
        # 步骤6: 展示如何使用 lm-eval 创建实例
        step6_create_hflm_instance()
        
    except Exception as e:
        print(f"\n错误: {e}")
        print("这可能是因为:")
        print("1. 网络连接问题（无法下载模型）")
        print("2. 模型名称不存在")
        print("3. 缺少必要的依赖")
        import traceback
        traceback.print_exc()
    
    print("\n" + "=" * 80)
    print("演示完成！")
    print("=" * 80)
    print("\n关键要点:")
    print("1. model_args 字符串被解析为字典")
    print("2. 根据 --model 参数获取对应的模型类（如 HFLM）")
    print("3. 使用 AutoConfig.from_pretrained() 加载 config.json")
    print("4. 使用 AutoTokenizer.from_pretrained() 加载 tokenizer")
    print("5. 使用 AutoModelForCausalLM.from_pretrained() 加载模型权重")
    print("\n详细说明请参考: LM_EVAL_MODEL_LOADING.md")


if __name__ == "__main__":
    main()