Точная настройка llama2 на cuda:1

Question

Когда я загружаю модель, я использую device_map для использования cuda:1, но кажется, что модель и обучение находятся на разных ядрах. Как мне правильно это сделать?

Код, работающий на Tesla T4 ниже:

      # load the base model in 4-bit quantization
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)
    
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    quantization_config=bnb_config,
    device_map={"": 1},
    trust_remote_code=True,
    use_auth_token=True,
)
base_model.config.use_cache = False
tokenizer = AutoTokenizer.from_pretrained(base_model_name, use_auth_token=True)
    
# add LoRA layers on top of the quantized base model
peft_config = LoraConfig(
    r=16,
    lora_alpha=64,
    lora_dropout=0.1,
    target_modules=["q_proj", "v_proj"],
    bias="none",
    task_type="CAUSAL_LM",
)
    
trainer = SFTTrainer(
    model=base_model,
    train_dataset=dataset,
    peft_config=peft_config,
    packing=True,
    max_seq_length=None,
    dataset_text_field="text",
    tokenizer=tokenizer,
    args=training_args,         # HF Trainer arguments
)
trainer.train()

Выдает ошибку:

ValueError: невозможно обучить модель, загруженную с точностью до 8 бит на устройстве, отличном от того, на котором вы тренируетесь. Убедитесь, что вы загрузили модель на правильное устройство, используя, например, device_map={'':torch.cuda.current_device()}, на котором вы тренируетесь.

Я следую этому руководству: https://huggingface.co/blog/dpo-trl .