Multiple gpu fine-tuning does not accelerate?

Question

I am experimenting using single or multiple GPUs for LLM fine-tuning by changing the CUDA_VISIBLE_DEVICES variable in the following cmd:

CUDA_VISIBLE_DEVICES=0,1 accelerate launch --multi_gpu finetuning_with_lora_HfArgumentParser.py \
--model_name "/root/123/local_model" \
--train_json_path "./train.json" \
--val_json_path "./val.json" \
--max_source_length 128 \
--max_target_length 256 \
--lora_rank 8 \
--lora_alpha 32 \
--output_dir "output" \
--logging_dir "logs" \
--num_train_epochs 10 \
--per_device_train_batch_size 1 \
--learning_rate 1e-4 \
--gradient_accumulation_steps 1024

However, I observed that total training time does not change regardless of CUDA_VISIBLE_DEVICES=0 or CUDA_VISIBLE_DEVICES=0,1. Does anyone know what's wrong with it?

Here is the training code:

import torch
from dataclasses import dataclass, field
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, TrainingArguments
from peft import LoraConfig, get_peft_model, TaskType
from qa_dataset import QADataset
from tqdm import tqdm
import time, sys

@dataclass
class CustomTrainingArguments: 
    model_name: str = field(
        default="Qwen/Qwen2-1.5B-Instruct",
        metadata={"help": "Pre-trained model for finetuning"}
    )
    train_json_path: str = field(
        default="./train.json",
        metadata={"help": "Path of training json data"}
    )
    val_json_path: str = field(
        default="./val.json",
        metadata={"help": "Path of validation json data"}
    )
    max_source_length: int = field(
        default=128,
        metadata={"help": "Maximum length of the input"}
    )
    max_target_length: int = field(
        default=256,
        metadata={"help": "Maximum length of the output"}
    )
    lora_rank: int = field(
        default=8,
        metadata={"help": "Inner dimension of the low-rank matrices to train"}
    )
    lora_alpha: int = field(
        default=32,
        metadata={"help": "Scaling factor for the low-rank matrices contribution"}
    )

def train_model(model, train_loader, val_loader, optimizer, gradient_accumulation_steps,
                device, num_epochs, model_output_dir, writer):
    batch_step = 0
    for epoch in range(int(num_epochs)):
        time1 = time.time()
        model.train()
        for index, data in enumerate(tqdm(train_loader, file=sys.stdout, desc="Train Epoch: " + str(epoch))):
            input_ids = data['input_ids'].to(device, dtype=torch.long)
            attention_mask = data['attention_mask'].to(device, dtype=torch.long)
            labels = data['labels'].to(device, dtype=torch.long)

            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels,
            )
            loss = outputs.loss

            loss.backward()

            if (index % gradient_accumulation_steps == 0 and index != 0) or index == len(train_loader) - 1:

                optimizer.step()

                optimizer.zero_grad()
                writer.add_scalar('Loss/train', loss, batch_step)
                batch_step += 1

            if index % 100 == 0 or index == len(train_loader) - 1:
                time2 = time.time()
                tqdm.write(
                    f"{index}, epoch: {epoch} -loss: {str(loss)} ; each step's time spent: {(str(float(time2 - time1) / float(index + 0.0001)))}")

        model.eval()
        val_loss = validate_model(model, val_loader, device)
        writer.add_scalar('Loss/val', val_loss, epoch)
        print(f"val loss: {val_loss} , epoch: {epoch}")
        print("Save Model To ", model_output_dir)
        model.save_pretrained(model_output_dir)


def validate_model(model, val_loader, device):
    running_loss = 0.0
    with torch.no_grad():
        for _, data in enumerate(tqdm(val_loader, file=sys.stdout, desc="Validation Data")):
            input_ids = data['input_ids'].to(device, dtype=torch.long)
            attention_mask = data['attention_mask'].to(device, dtype=torch.long)
            labels = data['labels'].to(device, dtype=torch.long)
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels,
            )
            loss = outputs.loss
            running_loss += loss.item()
    return running_loss / len(val_loader)


def main():

    parser = HfArgumentParser((TrainingArguments, CustomTrainingArguments))
    training_args, custom_args = parser.parse_args_into_dataclasses()


    # model_name = "Qwen/Qwen2-1.5B-Instruct"
    model_name = custom_args.model_name

    # train_json_path = "./train.json"
    train_json_path = custom_args.train_json_path

    # val_json_path = "./val.json"
    val_json_path = custom_args.val_json_path
    # max_source_length = 128
    max_source_length = custom_args.max_source_length
    # max_target_length = 256
    max_target_length = custom_args.max_target_length
    # epochs = 10
    epochs = training_args.num_train_epochs
    # batch_size = 1
    batch_size = training_args.per_device_train_batch_size
    # lr = 1e-4
    lr = training_args.learning_rate
    # gradient_accumulation_steps = 16
    gradient_accumulation_steps = training_args.gradient_accumulation_steps
    # lora_rank = 8
    lora_rank = custom_args.lora_rank
    # lora_alpha = 32
    lora_alpha = custom_args.lora_alpha
    # model_output_dir = "output"
    model_output_dir = training_args.output_dir
    # logs_dir = "logs"
    logs_dir = training_args.logging_dir

    # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    device = "cuda:{}".format(training_args.local_rank)
    print('------------------------')
    print(training_args.local_rank)
    print(device)

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
    # setup peft
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        inference_mode=False,
        r=lora_rank,
        lora_alpha=lora_alpha,
        lora_dropout=0.1
    )
    model = get_peft_model(model, peft_config)
    model.is_parallelizable = True
    model.model_parallel = True
    model.print_trainable_parameters()
    print("Start Load Train Data...")
    train_params = {
        "batch_size": batch_size,
        "shuffle": True,
        "num_workers": 0,
    }
    
    training_set = QADataset(train_json_path, tokenizer, max_source_length, max_target_length)
    training_loader = DataLoader(training_set, **train_params)
    print("Start Load Validation Data...")
    val_params = {
        "batch_size": batch_size,
        "shuffle": False,
        "num_workers": 0,
    }
    val_set = QADataset(val_json_path, tokenizer, max_source_length, max_target_length)
    val_loader = DataLoader(val_set, **val_params)
    
    writer = SummaryWriter(logs_dir)

    optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
    model = model.to(device)

    print("Start Training...")
    train_model(
        model=model,
        train_loader=training_loader,
        val_loader=val_loader,
        optimizer=optimizer,
        gradient_accumulation_steps=gradient_accumulation_steps,
        device=device,
        num_epochs=epochs,
        model_output_dir=model_output_dir,
        writer=writer
    )
    writer.close()

if __name__ == '__main__':
    main()