Как загрузить настроенную модель (объединенные веса) в Colab?

Question

Я настроил модель llama2. Перезагрузил базовую модель и объединил веса LoRA. Я снова сохранил эту наконец загруженную модель и теперь собираюсь ее запустить.

      base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    low_cpu_mem_usage=True,
    return_dict=True,
    torch_dtype=torch.float16,
    device_map=device_map,
)
model = PeftModel.from_pretrained(base_model, new_model)
model = model.merge_and_unload()
model.save_pretrained(...path/to/model)

Теперь я хотел бы использовать модель по пути/к/модели, используя следующий код

      
model_config = transformers.AutoConfig.from_pretrained(
    model_id,
    use_auth_token=hf_auth
)

model = transformers.AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    config=model_config,
    device_map='auto',
    offload_folder="offload",
    torch_dtype=float16,
    use_auth_token=hf_auth,
    offload_state_dict = True,
)
model.eval()

Моя цель сохранить объединенную модель — устранить зависимость от base_model.

проблема

Запуская модель в Colab, я вижу, что графический процессор не используется, а используется только процессор. Это приводит к сбою среды выполнения. Я хотел бы знать, что заставляет графический процессор не использоваться?