deepspeedai / deepspeedai/DeepSpeed

Issue with LoRA Tuning on llama3-70b using PEFT and TRL's SFTTrainer

Open
#5,719 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

training
Dominant language
Python
Stars
43.1k
Forks
5k
Avg merge
4d 15h
Merged PRs (30d)
112

Description

We are attempting to perform LoRA tuning on llama3-70b using PEFT with TRL's SFTTrainer. We are using 8 H100 GPUs and distributed training with ZeRO-stage3, but we encounter an error. Could you please provide any solutions?

Here is the error message:

Loading checkpoint shards:  77%|███████▋  | 23/30 [00:58<00:25,  3.68s/it]
Loading checkpoint shards:  80%|████████  | 24/30 [00:59<00:17,  2.87s/it]
Loading checkpoint shards:  83%|████████▎ | 25/30 [01:00<00:14,  2.89s/it]
Loading checkpoint shards:  80%|████████  | 24/30 [01:00<00:18,  3.12s/it]
Loading checkpoint shards:  83%|████████▎ | 25/30 [01:01<00:12,  2.47s/it]
Loading checkpoint shards:  87%|████████▋ | 26/30 [01:01<00:09,  2.49s/it]
Loading checkpoint shards:  83%|████████▎ | 25/30 [01:02<00:13,  2.70s/it]
Loading checkpoint shards:  87%|████████▋ | 26/30 [01:02<00:08,  2.18s/it]
Loading checkpoint shards:  90%|█████████ | 27/30 [01:03<00:06,  2.20s/it]
Loading checkpoint shards:  87%|████████▋ | 26/30 [01:03<00:09,  2.36s/it]
Loading checkpoint shards:  90%|█████████ | 27/30 [01:04<00:05,  1.99s/it]
Loading checkpoint shards:  93%|█████████▎| 28/30 [01:04<00:03,  1.99s/it]
Loading checkpoint shards:  90%|█████████ | 27/30 [01:05<00:06,  2.10s/it]
Loading checkpoint shards:  93%|█████████▎| 28/30 [01:05<00:03,  1.85s/it]
Loading checkpoint shards:  97%|█████████▋| 29/30 [01:06<00:01,  1.83s/it]
Loading checkpoint shards:  93%|█████████▎| 28/30 [01:06<00:03,  1.84s/it]
Loading checkpoint shards:  97%|█████████▋| 29/30 [01:06<00:01,  1.65s/it]
Loading checkpoint shards: 100%|██████████| 30/30 [01:06<00:00,  1.50s/it]
Loading checkpoint shards: 100%|██████████| 30/30 [01:06<00:00,  2.23s/it]

Loading checkpoint shards: 100%|██████████| 30/30 [01:07<00:00,  1.37s/it]
Loading checkpoint shards: 100%|██████████| 30/30 [01:07<00:00,  2.26s/it]
[WARNING|logging.py:314] 2024-07-02 18:12:30,312 >> Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.

Loading checkpoint shards:  97%|█████████▋| 29/30 [01:07<00:01,  1.59s/it][WARNING|logging.py:314] 2024-07-02 18:12:30,672 >> Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.

Loading checkpoint shards: 100%|██████████| 30/30 [01:08<00:00,  1.26s/it]
Loading checkpoint shards: 100%|██████████| 30/30 [01:08<00:00,  2.27s/it]
[WARNING|logging.py:314] 2024-07-02 18:12:31,194 >> Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
/home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/trl/trainer/utils.py:116: UserWarning: The pad_token_id and eos_token_id values of this tokenizer are identical. If you are planning for multi-turn training, it can result in the model continuously generating questions and answers without eos token. To avoid this, set the pad_token_id to a different value.
  warnings.warn(
/home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/trl/trainer/utils.py:116: UserWarning: The pad_token_id and eos_token_id values of this tokenizer are identical. If you are planning for multi-turn training, it can result in the model continuously generating questions and answers without eos token. To avoid this, set the pad_token_id to a different value.
  warnings.warn(
/home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/trl/trainer/utils.py:116: UserWarning: The pad_token_id and eos_token_id values of this tokenizer are identical. If you are planning for multi-turn training, it can result in the model continuously generating questions and answers without eos token. To avoid this, set the pad_token_id to a different value.
  warnings.warn(
[2024-07-02 18:12:37,303] [INFO] [launch.py:319:sigkill_handler] Killing subprocess 14959
[2024-07-02 18:12:37,304] [INFO] [launch.py:319:sigkill_handler] Killing subprocess 14960
Traceback (most recent call last):
  File "/work/scripts/train_py/run_clm_sft_update.py", line 686, in <module>
    main()
  File "/work/scripts/train_py/run_clm_sft_update.py", line 609, in main
    trainer = SFTTrainer(
  File "/home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/trl/trainer/sft_trainer.py", line 278, in __init__
    with PartialState().local_main_process_first():
  File "/home/user1/.pyenv/versions/3.10.14/lib/python3.10/contextlib.py", line 135, in __enter__
    return next(self.gen)
  File "/home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/accelerate/state.py", line 520, in local_main_process_first
    yield from self._goes_first(self.is_local_main_process)
  File "/home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/accelerate/state.py", line 384, in _goes_first
    self.wait_for_everyone()
  File "/home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/accelerate/state.py", line 378, in wait_for_everyone
    torch.distributed.barrier()
  File "/home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/distributed/c10d_logger.py", line 72, in wrapper
    return func(*args, **kwargs)
  File "/home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 3439, in barrier
    work = default_pg.barrier(opts=opts)
torch.distributed.DistBackendError: [3] is setting up NCCL communicator and retrieving ncclUniqueId from [0] via c10d key-value store by key '0', but store->get('0') got error: Connection reset by peer
Exception raised from recvBytes at ../torch/csrc/distributed/c10d/Utils.hpp:670 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7fe5eecf4d87 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libc10.so)
frame #1: <unknown function> + 0x5894fde (0x7fe5db5f0fde in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #2: c10d::TCPStore::doWait(c10::ArrayRef<std::string>, std::chrono::duration<long, std::ratio<1l, 1000l> >) + 0x360 (0x7fe5db5eb7f0 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #3: c10d::TCPStore::doGet(std::string const&) + 0x32 (0x7fe5db5ebb32 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #4: c10d::TCPStore::get(std::string const&) + 0xa1 (0x7fe5db5ec961 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #5: c10d::PrefixStore::get(std::string const&) + 0x31 (0x7fe5db5a1dd1 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #6: c10d::PrefixStore::get(std::string const&) + 0x31 (0x7fe5db5a1dd1 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #7: c10d::PrefixStore::get(std::string const&) + 0x31 (0x7fe5db5a1dd1 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #8: c10d::ProcessGroupNCCL::broadcastUniqueNCCLID(ncclUniqueId*, bool, std::string const&, int) + 0xa9 (0x7fe5a47dfc69 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #9: c10d::ProcessGroupNCCL::getNCCLComm(std::string const&, std::vector<c10::Device, std::allocator<c10::Device> > const&, c10d::OpType, int, bool) + 0x22b (0x7fe5a47e6c5b in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #10: <unknown function> + 0x10ad03d (0x7fe5a47f003d in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #11: c10d::ProcessGroupNCCL::allreduce_impl(std::vector<at::Tensor, std::allocator<at::Tensor> >&, c10d::AllreduceOptions const&) + 0x21 (0x7fe5a47f18e1 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #12: c10d::ProcessGroupNCCL::allreduce(std::vector<at::Tensor, std::allocator<at::Tensor> >&, c10d::AllreduceOptions const&) + 0x3bf (0x7fe5a47f38ff in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #13: c10d::ProcessGroupNCCL::barrier(c10d::BarrierOptions const&) + 0xb0e (0x7fe5a4802d4e in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #14: <unknown function> + 0x5838872 (0x7fe5db594872 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #15: <unknown function> + 0x5843590 (0x7fe5db59f590 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #16: <unknown function> + 0x5843695 (0x7fe5db59f695 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #17: <unknown function> + 0x4e8937c (0x7fe5dabe537c in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #18: <unknown function> + 0x1a08a38 (0x7fe5d7764a38 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #19: <unknown function> + 0x584cca4 (0x7fe5db5a8ca4 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #20: <unknown function> + 0x584da55 (0x7fe5db5a9a55 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #21: <unknown function> + 0xc93e88 (0x7fe5ede1ee88 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_python.so)
frame #22: <unknown function> + 0x413ef4 (0x7fe5ed59eef4 in /home/user1/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/lib/libtorch_python.so)
<omitting python frames>
frame #58: <unknown function> + 0x29d90 (0x7fe5ef964d90 in /usr/lib/x86_64-linux-gnu/libc.so.6)
frame #59: __libc_start_main + 0x80 (0x7fe5ef964e40 in /usr/lib/x86_64-linux-gnu/libc.so.6)
frame #60: _start + 0x25 (0x55b3be923095 in /home/user1/.pyenv/versions/3.10.14/bin/python3.10)
. This may indicate a possible application crash on rank 0 or a network set up issue.

Additionally, it's very strange because it worked correctly during a previous test run. Below is the log from that run. We haven't changed the code since then, but now we are encountering a new error.

One point of concern is that in the successful run log, there is a message:

[INFO|modeling_utils.py:3363] 2024-07-01 15:02:46,215 >> Detected DeepSpeed ZeRO-3: activating zero.init() for this model

before loading the model. However, this message is missing in the current log, and it seems the model is loaded into CPU memory first (previously, it was loaded directly into GPU memory).

Training Python Scripts:

import logging
import os
from contextlib import nullcontext

TRL_USE_RICH = os.environ.get("TRL_USE_RICH", False)

from trl.commands.cli_utils import init_zero_verbose, SFTScriptArguments, TrlParser

if TRL_USE_RICH:
    init_zero_verbose()
    FORMAT = "%(message)s"

    from rich.console import Console
    from rich.logging import RichHandler

import torch
from datasets import load_dataset

from tqdm.rich import tqdm
from transformers import AutoTokenizer

from trl import (
    ModelConfig,
    RichProgressCallback,
    SFTConfig,
    SFTTrainer,
    get_peft_config,
    get_quantization_config,
    get_kbit_device_map,
)

tqdm.pandas()

if TRL_USE_RICH:
    logging.basicConfig(format=FORMAT, datefmt="[%X]", handlers=[RichHandler()], level=logging.INFO)


if __name__ == "__main__":
    parser = TrlParser((SFTScriptArguments, SFTConfig, ModelConfig))
    args, training_args, model_config = parser.parse_args_and_config()

    # Force use our print callback
    if TRL_USE_RICH:
        training_args.disable_tqdm = True
        console = Console()

    ################
    # Model & Tokenizer
    ################
    torch_dtype = (
        model_config.torch_dtype
        if model_config.torch_dtype in ["auto", None]
        else getattr(torch, model_config.torch_dtype)
    )
    quantization_config = get_quantization_config(model_config)
    model_kwargs = dict(
        revision=model_config.model_revision,
        trust_remote_code=model_config.trust_remote_code,
        attn_implementation=model_config.attn_implementation,
        torch_dtype=torch_dtype,
        use_cache=False if training_args.gradient_checkpointing else True,
        device_map=get_kbit_device_map() if quantization_config is not None else None,
        quantization_config=quantization_config,
    )
    tokenizer = AutoTokenizer.from_pretrained(model_config.model_name_or_path, use_fast=True)
    tokenizer.pad_token = tokenizer.eos_token

    ################
    # Dataset
    ################
    raw_datasets = load_dataset(args.dataset_name)

    train_dataset = raw_datasets[args.dataset_train_split]
    eval_dataset = raw_datasets[args.dataset_test_split]

    ################
    # Optional rich context managers
    ###############
    init_context = nullcontext() if not TRL_USE_RICH else console.status("[bold green]Initializing the SFTTrainer...")
    save_context = (
        nullcontext()
        if not TRL_USE_RICH
        else console.status(f"[bold green]Training completed! Saving the model to {training_args.output_dir}")
    )

    ################
    # Training
    ################
    with init_context:
        trainer = SFTTrainer(
            model=model_config.model_name_or_path,
            model_init_kwargs=model_kwargs,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            tokenizer=tokenizer,
            peft_config=get_peft_config(model_config),
            callbacks=[RichProgressCallback] if TRL_USE_RICH else None,
        )

    trainer.train()

    with save_context:
        trainer.save_model(training_args.output_dir)

Training ShellScripts:

export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"

time \
deepspeed \
sft.py \
--deepspeed ds_config_zero3.json \
--dataset_dir mytest \
--model_name_or_path meta-llama/Meta-Llama-3-70B-Instruct \
--tokenizer_name meta-llama/Meta-Llama-3-70B-Instruct \
--num_train_epochs 5 \
--do_train \
--do_eval \
--bf16 \
--output_dir ./lora-test \
--per_device_train_batch_size 1 \
--per_device_eval_batch_size 1 \
--gradient_accumulation_steps 16 \
--learning_rate=5e-6 \
--lr_scheduler_type "constant" \
--warmup_ratio 0.03 \
--logging_steps 1 \
--evaluation_strategy steps \
--evaluation_steps 100 \
--save_strategy epoch \
--overwrite_output_dir \
--gradient_checkpointing \
--use_peft True \
--lora_r 16 \
--ddp_timeout 72000 \
--lora_alpha 32 \
--lora_dropout 0.05 \
--lora_target_modules q_proj v_proj k_proj o_proj gate_proj down_proj up_proj \

DeepSpeed Config:

{
    "bf16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
 
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": 5e7,
        "stage3_prefetch_bucket_size": 5e7,
        "stage3_param_persistence_threshold": 0,
        "stage3_max_live_parameters": 1e8,
        "stage3_max_reuse_distance": 1e8, 
        "sub_group_size": 5e7,
        "stage3_gather_fp16_weights_on_model_save": true
    },
 
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with /work/scripts/train_py/run_clm_sft_update.py around the SFTTrainer construction at line 609, then trace the reported failure through accelerate/state.py and torch.distributed.barrier(). Compare the successful and failing logs, including DeepSpeed ZeRO-3 initialization, and identify whether rank 0 or the network setup fails before considering a fix.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.