simplecloud/VidChain-exercise

1import os2import shutil3 4from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig5import torch6from vtimellm.model import *7from peft import PeftModel8 9def load_lora(model, lora_path):10    non_lora_trainables_path = os.path.join(lora_path, 'non_lora_trainables.bin')11    if os.path.exists(non_lora_trainables_path):12        non_lora_trainables = torch.load(non_lora_trainables_path, map_location='cpu')13        non_lora_trainables = {(k[11:] if k.startswith('base_model.') else k): v for k, v in non_lora_trainables.items()}14        if any(k.startswith('model.model.') for k in non_lora_trainables):15            non_lora_trainables = {(k[6:] if k.startswith('model.') else k): v for k, v in non_lora_trainables.items()}16        model.load_state_dict(non_lora_trainables, strict=False)17    print('Loading LoRA weights...')18    model = PeftModel.from_pretrained(model, lora_path)19    return model20 21def load_pretrained_model(args, stage2=None, stage3=None, stage4=None, stage5=None):22    """23    Load VTimeLLM model with proper GPU device handling24    25    FIXED VERSION: This function now properly handles GPU device selection26    to prevent multiple GPU detection issues.27    """28    kwargs = {'torch_dtype': torch.float16}29 30    # model_path = os.path.expanduser(args.model_path)31    model_base = args.model_base32 33    # FIX: Set up single GPU environment before model loading34    if torch.cuda.is_available():35        # Ensure we're using the correct GPU device36        current_device = torch.cuda.current_device()37        print(f'Using GPU device: {current_device}')38        print(f'GPU name: {torch.cuda.get_device_name(current_device)}')39        40        # Set device for all operations41        device = f'cuda:{current_device}'42    else:43        device = 'cpu'44        print('No CUDA available, using CPU')45 46    # lora_cfg_pretrained = AutoConfig.from_pretrained(model_path)47    print('Loading VTimeLLM from base model...')48    if 'chatglm' in model_base:49        tokenizer = AutoTokenizer.from_pretrained(model_base, trust_remote_code=True)50        model = VTimeLLMChatGLMForCausalLM.from_pretrained(model_base)51    else:52        tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)53        model = VTimeLLMLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, **kwargs)54        token_num, tokem_dim = model.lm_head.out_features, model.lm_head.in_features55        if model.lm_head.weight.shape[0] != token_num:56            model.lm_head.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))57            model.model.embed_tokens.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))58 59    # FIX: Move model to GPU with explicit device selection60    if torch.cuda.is_available():61        model = model.to(device)62        print(f'Model moved to {device}')63    else:64        print('Model loaded on CPU')65    66    model.get_model().initialize_vision_modules(args)67 68    if stage2 is not None and stage2 != "":69        print('Loading stage2 weights...')70        model = load_lora(model, stage2)71        print('Merging stage2 weights...')72        model = model.merge_and_unload()73 74        if stage3 is not None and stage3 != "" :75            print('Loading stage3 weights...')76            model = load_lora(model, stage3)77            print('Merging stage3 weights...')78            model = model.merge_and_unload()79 80        if stage4 is not None and stage4 != "":81            print('Loading stage4 weights...')82            model = load_lora(model, stage4)83            print('Merging stage4 weights...')84            model = model.merge_and_unload()      85 86        if stage5 is not None and stage5 != "":87            print('Loading stage5 weights...')88            model = load_lora(model, stage5)89            print('Merging stage5 weights...')90            model = model.merge_and_unload()91 92    if hasattr(model.config, "max_sequence_length"):93        context_len = model.config.max_sequence_length94    else:95        context_len = 204896 97    return tokenizer, model, context_len98 99def load_pretrained_model_single_gpu(args, stage2=None, stage3=None, stage4=None, stage5=None):100    """101    Load VTimeLLM model with forced single GPU usage102    103    This function ensures only one GPU is used by setting environment variables104    and explicitly managing device placement.105    """106    # Force single GPU usage107    os.environ['CUDA_VISIBLE_DEVICES'] = '0'108    os.environ['OMPI_COMM_WORLD_SIZE'] = '1'109    110    if torch.cuda.is_available():111        torch.cuda.set_device(0)112        print(f'Forced single GPU usage: {torch.cuda.get_device_name(0)}')113    114    return load_pretrained_model(args, stage2, stage3, stage4, stage5)115