[BUG] train_batch_size is not equal to micro_batch_per_gpu * gradient_acc_step * world_size

[vecorro]

Describe the bug

When using DeepSpeed 0.10.0 (or version > 0.8.2) with Ray 2.5.1 I get the following error when trying to run a job on 3 Reay workers:

AssertionError: Check batch related parameters. train_batch_size is not equal to micro_batch_per_gpu * gradient_acc_step * world_size 9 != 1 * 3 * 1

To Reproduce
Steps to reproduce the behavior:

Run the following script on a Ray cluster with 3 nodes, each hosting 1 NVIDIA GPU A100 (40GB)
ray job submit --address <head-IP> --working-dir ./src/ --runtime-env-json='{"pip": ["torch==2.0.1","transformers==4.30.2","deepspeed==0.10.0", "accelerate==0.21.0", "peft==0.4.0", "bitsandbytes==0.39.1","datasets==2.13.0", "einops==0.6.1"]}' -- python finetune.py --num-workers=3 --model="tiiuae/falcon-7b"

Expected behavior
I expected this release to include a fix to #3228

ds_report output
I cannot run ds_repost as I'm using a Ray cluster, but the issue seems to be the same as #3228, although I'm using Xero stage 3.

DS config

deepspeed_cfg = {'fp16': {'enabled': 'auto',
                     'loss_scale': 0,
                     'loss_scale_window': 1000,
                     'initial_scale_power': 16,
                     'hysteresis': 2,
                     'min_loss_scale': 1},
            'optimizer': {'type': 'AdamW',
                          'params': {'lr': 'auto',
                                     'betas': 'auto',
                                     'eps': 'auto',
                                     'weight_decay': 'auto'}},
            'scheduler': {'type': 'WarmupLR',
                          'params': {'warmup_min_lr': 'auto',
                                     'warmup_max_lr': 'auto',
                                     'warmup_num_steps': 'auto',
                                     'warmup_type': 'linear'}},
            'zero_optimization': {'stage': 3,
                                  'offload_optimizer': {'device': 'cpu', 'pin_memory': False},
                                  'offload_param': {'device': 'cpu', 'pin_memory': False},
                                  'overlap_comm': False,
                                  'contiguous_gradients': True,
                                  'reduce_bucket_size': 'auto',
                                  'stage3_prefetch_bucket_size': 'auto',
                                  'stage3_param_persistence_threshold': 'auto',
                                  'sub_group_size': 1000000000.0,
                                  'stage3_max_live_parameters': 1000000000.0,
                                  'stage3_max_reuse_distance': 1000000000.0,
                                  'stage3_gather_16bit_weights_on_model_save': 'auto'},
            'gradient_accumulation_steps': 'auto',
            'gradient_clipping': False,
            'steps_per_print': 10,
            'train_batch_size': 'auto',
            'train_micro_batch_size_per_gpu': 'auto',
            'wall_clock_breakdown': False}

If applicable, add screenshots to help explain your problem.

System info (please complete the following information):

• OS: Ray image rayproject/ray:nightly-py39-cu118
• GPU count and types: 3 Ray 2.5.1 workers each with 1 x NVIDIA A100 (40 GB) GPU
• Python version: 3.9
• CUDA 11.8
• Ray working env: --runtime-env-json='{"pip": ["torch==2.0.1","transformers==4.30.2","deepspeed==0.10.0", "accelerate==0.21.0", "peft==0.4.0", "bitsandbytes==0.39.1","datasets==2.13.0", "einops==0.6.1"]}'

Additional context
SCRIPT

import os

os.environ['LD_LIBRARY_PATH'] = f"/usr/local/cuda/lib64:{os.environ['LD_LIBRARY_PATH']}"
os.environ['PATH'] = f"/usr/local/cuda/bin:{os.environ['PATH']}"
print(f"LD_LIBRARY_PATH: {os.environ['LD_LIBRARY_PATH']}")
print(f"PATH: {os.environ['PATH']}")

import ray
import argparse
import torch
import json
import pandas as pd
from pprint import pprint
from datasets import load_dataset
from ray.data.preprocessors import BatchMapper
from ray.train.huggingface import TransformersTrainer
from ray.air.config import ScalingConfig
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import TrainingArguments
from transformers import Trainer
from transformers import DataCollatorForLanguageModeling
from peft import prepare_model_for_int8_training
from peft import LoraConfig, get_peft_model
from ray.air import session

# LoRA Configuration
LORA_R = 8
LORA_ALPHA = 16
LORA_DROPOUT = 0.05
# LORA_TARGET_MODULES = ['query_key_value']
LORA_TARGET_MODULES = [
        "query_key_value",
        "dense",
        "dense_h_to_4h",
        "dense_4h_to_h",
]


def load_model(model_id, lora_config):
    # Load Model
    model = AutoModelForCausalLM.from_pretrained(
            model_id,
            load_in_8bit=True,
            torch_dtype=torch.float16,
            trust_remote_code=True,
            #device_map="auto"
    )
    model.config.use_cache = False
    model = prepare_model_for_int8_training(model)
    model = get_peft_model(model, lora_config)
    return model


def trainer_init_per_worker(train_dataset, eval_dataset=None, **config):
    print(f"Is CUDA available? {torch.cuda.is_available()}")

    # Use the actual number of CPUs assigned by Ray
    os.environ["OMP_NUM_THREADS"] = str(
            session.get_trial_resources().bundles[-1].get("CPU", 1)
    )

    # Enable tf32 for better performance
    torch.backends.cuda.matmul.allow_tf32 = True

    # Load tokenizer
    model_id = config.get('model')
    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
    tokenizer.pad_token = tokenizer.eos_token

    # Apply LoRA
    lora_config = LoraConfig(
            r=config.get('lora').get('r'),
            lora_alpha=config.get('lora').get('alpha'),
            target_modules=config.get('lora').get('target_modules'),
            lora_dropout=config.get('lora').get('dropout'),
            bias="none",
            task_type="CAUSAL_LM",
    )

    # Training Configuration
    training_args = TrainingArguments(
            per_device_train_batch_size=1,
            # auto_find_batch_size=True,
            gradient_accumulation_steps=3,
            num_train_epochs=1,
            learning_rate=2e-4,
            fp16=False,
            bf16=False,
            save_total_limit=4,
            logging_steps=5,
            save_strategy='steps',
            weight_decay=0,
            push_to_hub=False,
            disable_tqdm=True,
            no_cuda=not config.get('platform').get('use_gpu'),
            gradient_checkpointing=True,
            output_dir="./outputs_ray",
            ddp_find_unused_parameters=False,
            deepspeed=config.get('platform').get('deepspeed'),
            max_steps=40,
    )

    # Data Collator
    data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

    # Configure the Trainer
    trainer = Trainer(
            model=load_model(model_id, lora_config),
            train_dataset=train_dataset,
            args=training_args,
            data_collator=data_collator
    )

    return trainer


def prepare_dataset(path, model_id):
    dataset = load_dataset(path, split="train")

    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
    tokenizer.pad_token = tokenizer.eos_token

    dataset_prompts = {}
    dataset_prompts['text'] = []

    def generate_prompt(data_point):
        return f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately 
    completes the request.  # noqa: E501
    ### Instruction:
    {data_point["instruction"]}
    ### Input:
    {data_point["input"]}
    ### Response:
    {data_point["output"]}"""

    for data_point in dataset:
        prompt = generate_prompt(data_point)
        dataset_prompts['text'].append(prompt)

    # Transform to Ray dataset format
    dataset_prompts_df = pd.DataFrame.from_dict(dataset_prompts)
    dataset_ray = ray.data.from_pandas(dataset_prompts_df)

    return dataset_ray


def prepare_batch_mapper(model_id):
    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
    tokenizer.pad_token = tokenizer.eos_token

    def preprocess_function(batch):
        ret = tokenizer(list(batch["text"]), padding=True, truncation=True, return_tensors="np")
        return dict(ret)

    batch_mapper = BatchMapper(preprocess_function, batch_format="pandas")

    return batch_mapper


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--num-workers", type=int, default=2, help="Sets number of workers for training.")
    parser.add_argument("--use-cpu", action="store_true", default=False, help="Enables CPU training")
    parser.add_argument("--no-deepspeed", action="store_true", default=False, help="Disables DeepSpeed strategy")
    parser.add_argument('--model', action='store', type=str, default="tiiuae/falcon-7b",
                        help='Model from HuggingFace to use')
    parser.add_argument('--data', action='store', type=str, default="yahma/alpaca-cleaned",
                        help='Path of the data to use for finetuning the model')

    args, _ = parser.parse_known_args()

    # Get DeepSpeed config
    if not args.no_deepspeed:
        deepspeed = {
                "fp16": {
                        "enabled": "auto",
                        "loss_scale": 0,
                        "loss_scale_window": 1000,
                        "initial_scale_power": 16,
                        "hysteresis": 2,
                        "min_loss_scale": 1
                },

                "optimizer": {
                        "type": "AdamW",
                        "params": {
                                "lr": "auto",
                                "betas": "auto",
                                "eps": "auto",
                                "weight_decay": "auto"
                        }
                },

                "scheduler": {
                        "type": "WarmupLR",
                        "params": {
                                "warmup_min_lr": "auto",
                                "warmup_max_lr": "auto",
                                "warmup_num_steps": "auto"
                        }
                },

                "zero_optimization": {
                        "stage": 3,
                        "offload_optimizer": {
                                "device": "cpu",
                                "pin_memory": False
                        },
                        "offload_param": {
                                "device": "cpu",
                                "pin_memory": False
                        },
                        "overlap_comm": False,
                        "contiguous_gradients": True,
                        "sub_group_size": 1e9,
                        "reduce_bucket_size": "auto",
                        "stage3_prefetch_bucket_size": "auto",
                        "stage3_param_persistence_threshold": "auto",
                        "stage3_max_live_parameters": 1e9,
                        "stage3_max_reuse_distance": 1e9,
                        "stage3_gather_16bit_weights_on_model_save": True
                },

                "gradient_accumulation_steps": "auto",
                "gradient_clipping": False,
                "steps_per_print": 10,
                "train_batch_size": "auto",
                "train_micro_batch_size_per_gpu": "auto",
                "wall_clock_breakdown": False,
        }

    # Init Ray cluster
    ray.init(address="auto")
    print(f" Ray CLuster resources:\n {ray.cluster_resources()}")

    # Prepare Ray dataset and batch mapper
    dataset = prepare_dataset(args.data, args.model)
    batch_mapper = prepare_batch_mapper(args.model)

    # Trainer init config
    trainer_init_config = {"model": args.model,
                           "lora": {"r": LORA_R,
                                    "alpha": LORA_ALPHA,
                                    "dropout": LORA_DROPOUT,
                                    "target_modules": LORA_TARGET_MODULES
                                    },
                           "platform": {"use_gpu": not args.use_cpu,
                                        "deepspeed": deepspeed if not args.no_deepspeed else None}
                           }

    trainer = TransformersTrainer(
            trainer_init_per_worker=trainer_init_per_worker,
            trainer_init_config=trainer_init_config,
            scaling_config=ScalingConfig(num_workers=args.num_workers, use_gpu=not args.use_cpu),
            datasets={
                    "train": dataset
            },
            preprocessor=batch_mapper,
    )

    # Launch the training on the cluster
    result = trainer.fit()

[kouroshHakha]

Hey @vecorro,

From ray team here. I have seen this issue reported from other folks when they use the TransformerTrainer. I have not yet got a chance to look into the that issue yet, but I got the Llama-2 fine-tuning working (for all params) with the following settings:

Cluster type: 4xg5.48xlarge = 32xA10Gs

Deepspeed config:

    deepspeed_config = {
        "fp16": {
            "enabled": False
        },
        "bf16": {
            "enabled": True
        },
        "scheduler": {
            "type": "WarmupLR",
            "params": {
                "warmup_min_lr": "auto",
                "warmup_max_lr": "auto",
                "warmup_num_steps": "auto"
            }
        },
        "optimizer": {
            "type": "AdamW",
            "params": {
            "lr": 1e-5,
            "betas": [0.99, 0.999],
            "eps": 1e-8,
            "weight_decay": 1e-6
            }
        },
        "zero_optimization": {
            "stage": 3,
            "offload_optimizer": {
                "device": "cpu",
                "pin_memory": False,
            },
            # "offload_param": {
            #     "device": "cpu",
            #     "pin_memory": True,
            # },
            "overlap_comm": True,
            "contiguous_gradients": True,
            "reduce_bucket_size": "auto",
            "stage3_prefetch_bucket_size": "auto",
            "stage3_param_persistence_threshold": "auto",
            "gather_16bit_weights_on_model_save": True,
            "round_robin_gradients": True,
        },
        "gradient_accumulation_steps": "auto",
        "gradient_clipping": "auto",
        "steps_per_print": 10,
        "train_batch_size": "auto",
        "train_micro_batch_size_per_gpu": batch_size,
        "wall_clock_breakdown": False,
    }

Cluster env config:

# See https://hub.docker.com/r/anyscale/ray for full list of
# available Ray, Python, and CUDA versions.
base_image: anyscale/ray:2.5.1-py39-cu116

env_vars: {}

debian_packages: []

python:
  pip_packages: []
  conda_packages: []

post_build_cmds: [
    pip install torch==2.0.0+cu117 torchvision==0.15.1+cu117 torchaudio==2.0.1 --index-url https://download.pytorch.org/whl/cu117,
    pip install flash-attn==1.0.3.post0,
    pip install git+https://github.com/EleutherAI/lm_dataformat.git@4eec05349977071bf67fc072290b95e31c8dd836,
    pip install deepspeed==0.10.0,
    pip install fairscale==0.4.13,
    pip install transformers==4.31.0 ,
    pip install dataset==1.6.2 ,
    pip install accelerate==0.21.0 ,
    pip install evaluate,
    pip install bitsandbytes ,
    pip install wandb,
    pip install pytorch-lightning ,
    pip install protobuf<3.21.0 ,
    pip install torchmetrics ,
    pip install lm_eval==0.3.0, 
    pip install tiktoken==0.1.2,
    pip install sentencepiece,
]

Script:

import evaluate
import torch
import argparse
import json
import math
import os
import functools
import time 
import tree
import pandas as pd
from pathlib import Path
import torch
import torch.nn as nn
from ray import tune

from accelerate import Accelerator, DeepSpeedPlugin
from accelerate.utils import DummyOptim, DummyScheduler, set_seed
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
)


import ray
from ray import air
from ray.air import session
from ray.train.torch import TorchTrainer
import ray.util.scheduling_strategies

from utils import get_checkpoint_and_refs_dir, get_mirror_link

OPTIM_BETAS = (0.9, 0.999)
OPTIM_EPS = 1e-8
OPTIM_WEIGHT_DECAY = 0.

MIRROR_LINK = "s3://llama-2-weights/"

def get_number_of_params(model:nn.Module):
    state_dict = model.state_dict()
    return sum(p.numel() for p in state_dict.values())

def collate_fn(batch, tokenizer, block_size, device):
    out_batch = tokenizer(
        list(batch["input"]), 
        padding="longest", 
        max_length=block_size, 
        truncation=True, 
        return_tensors="pt"    
    )
    out_batch["labels"] = out_batch["input_ids"].clone()
    
    out_batch = tree.map_structure(lambda x: x.to(device), out_batch)
    
    return out_batch

def get_pretrained_path(model_id: str):
    mirror_uri = get_mirror_link(model_id)    
    ckpt_path, _ = get_checkpoint_and_refs_dir(
        model_id=model_id,
        bucket_uri=mirror_uri
    )
    return ckpt_path


def get_tokenizer(model_name, special_tokens):

    pretrained_path = get_pretrained_path(model_name)
    tokenizer = AutoTokenizer.from_pretrained(pretrained_path)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.add_tokens(special_tokens, special_tokens=True)
    
    return tokenizer
    
    
def create_ray_dataset(path):
    # jsonl file
    with open(path, 'r') as json_file:
        items = [json.loads(x) for x in json_file]

    dataset = {"input": []}
    for item in items:
        assert set(item.keys()) == {"input"}
        dataset["input"].append(item["input"])
    
    df = pd.DataFrame.from_dict(dataset)
    return ray.data.from_pandas(df)

def evaluate(model, eval_dataset_len, eval_dataloader, accelerator, bsize):
    model.eval()
    losses = []
    for step, batch in enumerate(eval_dataloader):
        with torch.no_grad():
            outputs = model(**batch)

        loss = outputs.loss
        losses.append(accelerator.gather(loss.repeat(bsize)))

    losses = torch.cat(losses)
    losses = losses[: eval_dataset_len]
    try:
        eval_loss = torch.mean(losses)
        perplexity = math.exp(eval_loss)
    except OverflowError:
        perplexity = float("inf")
    return perplexity, eval_loss


def checkpoint_model(checkpoint_folder, ckpt_id, model, epoch, last_global_step, **kwargs):
    """Utility function for checkpointing model + optimizer dictionaries
    The main purpose for this is to be able to resume training from that instant again.
    """
    checkpoint_state_dict = {
        "epoch": epoch,
        "last_global_step": last_global_step,
    }
    # Add extra kwargs too
    checkpoint_state_dict.update(kwargs)
    
    # In here model will be a DeepspeedEngine object
    model.save_checkpoint(checkpoint_folder, ckpt_id, checkpoint_state_dict)
    status_msg = f"checkpointing: checkpoint_folder={checkpoint_folder}, ckpt_id={ckpt_id}"
    print(status_msg)

def training_function(kwargs: dict):
    os.environ["OMP_NUM_THREADS"] = str(
        session.get_trial_resources().bundles[-1].get("CPU", 1)
    )

    print("training_function called")

    config = kwargs["config"]
    args = argparse.Namespace(**kwargs["args"])
    special_tokens = kwargs.get("special_tokens", [])
    model_id = config["model_name"]

    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
    lr = config["lr"]
    num_epochs = int(config["num_epochs"])
    seed = int(config["seed"])
    batch_size = int(config["batch_size"])
    gradient_accumulation_steps = int(config["gradient_accumulation_steps"])
    as_test = config.get("as_test", False)

    deepspeed_config = {
        "fp16": {
            "enabled": False
        },
        "bf16": {
            "enabled": True
        },
        "scheduler": {
            "type": "WarmupLR",
            "params": {
                "warmup_min_lr": "auto",
                "warmup_max_lr": "auto",
                "warmup_num_steps": "auto"
            }
        },
        "optimizer": {
            "type": "AdamW",
            "params": {
            "lr": 1e-5,
            "betas": [0.99, 0.999],
            "eps": 1e-8,
            "weight_decay": 1e-6
            }
        },
        "zero_optimization": {
            "stage": 3,
            "offload_optimizer": {
                "device": "cpu",
                "pin_memory": False,
            },
            # "offload_param": {
            #     "device": "cpu",
            #     "pin_memory": True,
            # },
            "overlap_comm": True,
            "contiguous_gradients": True,
            "reduce_bucket_size": "auto",
            "stage3_prefetch_bucket_size": "auto",
            "stage3_param_persistence_threshold": "auto",
            "gather_16bit_weights_on_model_save": True,
            "round_robin_gradients": True,
        },
        "gradient_accumulation_steps": "auto",
        "gradient_clipping": "auto",
        "steps_per_print": 10,
        "train_batch_size": "auto",
        "train_micro_batch_size_per_gpu": batch_size,
        "wall_clock_breakdown": False,
    }
    ds_plugin = DeepSpeedPlugin(hf_ds_config=deepspeed_config)

    # Initialize accelerator
    accelerator = Accelerator(
        deepspeed_plugin=ds_plugin,
        gradient_accumulation_steps=gradient_accumulation_steps,
        mixed_precision=args.mx,
    )
    
    set_seed(seed)
    
    # train_ds is the local shard for this model
    train_ds = session.get_dataset_shard("train")
    valid_ds = session.get_dataset_shard("valid")
    
    train_ds_len = len(list(train_ds.iter_batches(batch_size=1)))
    valid_ds_len = len(list(valid_ds.iter_batches(batch_size=1)))
    
    tokenizer = get_tokenizer(model_name=args.model_name, special_tokens=special_tokens)
    collate_partial = functools.partial(
        collate_fn, 
        tokenizer=tokenizer,
        block_size=config["block_size"], 
        device=accelerator.device
    )
    
    ckpt_path = session.get_trial_dir()

    pretrained_path = get_pretrained_path(model_id)    
    print(f"Loading model from {pretrained_path} ...")
    s = time.time()
    model = AutoModelForCausalLM.from_pretrained(
        pretrained_path,
        trust_remote_code=True, 
        torch_dtype=torch.bfloat16,
    )
    print(f"Done loading model in {time.time() - s} seconds.")
    model.resize_token_embeddings(len(tokenizer))    
    print("Model initialized with pretrained weights. Training starting...")
    if not args.no_grad_ckpt:
        model.gradient_checkpointing_enable()

    optimizer_cls = (
        torch.optim.AdamW
        if accelerator.state.deepspeed_plugin is None
        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
        else DummyOptim
    )

    optimizer = optimizer_cls(
        model.parameters(), 
        lr=lr, 
        betas=OPTIM_BETAS, 
        weight_decay=OPTIM_WEIGHT_DECAY, 
        eps=OPTIM_EPS
    )

    # Instantiate scheduler
    # Creates Dummy Scheduler if `scheduler` was spcified in the config file else creates `args.lr_scheduler_type` Scheduler
    # get train and valid dataset lengths

    if (
        accelerator.state.deepspeed_plugin is None
        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
    ):
        lr_scheduler = get_linear_schedule_with_warmup(
            optimizer=optimizer,
            num_warmup_steps=100,
            num_training_steps=(train_ds_len * num_epochs) // gradient_accumulation_steps,
        )
    else:
        lr_scheduler = DummyScheduler(
            optimizer, total_num_steps=(train_ds_len * num_epochs) // gradient_accumulation_steps, warmup_num_steps=100
        )

    # Prepare everything
    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the prepare method.
    s = time.time()
    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
    print(f"Prepare done in {time.time() - s} seconds.")

    # Now we train the model
    if accelerator.is_main_process:
        print("Starting training ...")
        print("number of batches on main process", train_ds_len // batch_size)

    avg_fwd_time, avg_bwd_time, avg_opt_step_time = 0, 0, 0
    for epoch in range(num_epochs):
        s_epoch = time.time()
        model.train()
        loss_sum = torch.tensor(0.0).to(accelerator.device)
        
        train_dataloader = train_ds.iter_torch_batches(
            batch_size=batch_size,
            collate_fn=collate_partial,
        )
    
        for step, batch in enumerate(train_dataloader):
            
            # We could avoid this line since we set the accelerator with `device_placement=True`.
            with accelerator.accumulate(model):
                s_fwd = time.time()
                outputs = model(**batch)
                loss = outputs.loss
                loss_sum += loss
                e_fwd = time.time()
                avg_fwd_time += e_fwd - s_fwd
                s_bwd = time.time()
                accelerator.backward(loss)
                e_bwd = time.time()
                avg_bwd_time += e_bwd - s_bwd

                # if accelerator.sync_gradients:
                #     accelerator.clip_grad_norm_(model.parameters(), 1.0)

                s_opt_step = time.time()
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()
                e_opt_step = time.time()
                avg_opt_step_time += e_opt_step - s_opt_step
                
            if accelerator.is_main_process:
                accelerator.print(
                    f"[epoch {epoch} step {step}] loss: {loss.item()} step-time: {e_opt_step - s_fwd}"
                )
                
            # as long as this is not the last step report here
            if step != (train_ds_len // batch_size - 1) and not as_test:
                session.report(
                    {
                        "epoch": epoch,
                        "step": step,
                        "train_loss": loss.item(),
                        "eval_loss": None,
                        "perplexity": None,
                        "number of iterations": step + 1,
                        "Train time per epoch": None,
                        "Eval time per epoch": None,
                        "avg fwd time": avg_fwd_time / (step + 1),
                        "avg bwd time": avg_bwd_time / (step + 1),
                    },
                )
            
            if as_test:
                break
            

        e_epoch = time.time()
        accelerator.print("Train time per epoch: ", e_epoch - s_epoch)

        eval_s_epoch = time.time()        
        valid_dataloader = valid_ds.iter_torch_batches(
            batch_size=batch_size,
            collate_fn=collate_partial,
        )
        
        perplex, eloss = evaluate(
            model, 
            valid_ds_len, 
            valid_dataloader, 
            accelerator, 
            bsize=config["eval_batch_size"]
        )
        accelerator.print("Eval result loss", eloss)
        accelerator.print("Eval perplex", perplex)

        eval_e_epoch = time.time()
        accelerator.print("Eval time per epoch: ", eval_e_epoch - eval_s_epoch)
        accelerator.print("avg fwd time: ", avg_fwd_time / (step + 1))
        accelerator.print("avg bwd time: ", avg_bwd_time / (step + 1))
        accelerator.print("avg opt step time: ", avg_opt_step_time / (step + 1))

        accelerator.print(f"Saving the model in {ckpt_path}")
        
        with accelerator.main_process_first():
            ckpt_path_epoch = Path(ckpt_path) / f"epoch-{epoch}"
            print(f"Creating the checkpoint directory {ckpt_path_epoch} starting from main process first.")
            ckpt_path_epoch.mkdir(parents=True, exist_ok=True)
            
            
            print("Saving tokenizer and config in the checkpoint directory ...")
            tokenizer.save_pretrained(ckpt_path_epoch)

        print("Waiting for everyone to get to here ...")
        # accelerator.wait_for_everyone()
        
        # This checkpointing method makes deepspeed checkpoints on each node and then 
        # Ray Train will aggregate them to a central s3 bucket.
        # It should not be done on all processes (not just the Rank 0)
        # checkpoint_model(
        #     checkpoint_folder=ckpt_path_epoch, 
        #     ckpt_id=epoch, 
        #     model=model, 
        #     epoch=epoch, 
        #     last_global_step=step
        # )
        
        # unwrapped_model = accelerator.unwrap_model(model)
        # unwrapped_model.save_pretrained(
        #     ckpt_path_epoch,
        #     is_main_process=accelerator.is_main_process,
        #     save_function=accelerator.save,
        #     safe_serialization=True,
        #     state_dict=accelerator.get_state_dict(model),
        # )
        
        # Note: After the following call, in the case of remote storage, the checkpoint 
        # directoy will get synced to the remote storage and then deleted from the 
        # local directory. This will open up local disk. 
        # session.report(
        #     {
        #         "epoch": epoch,
        #         "iteration": step,
        #         "train_loss": loss_sum.item() / (step + 1),
        #         "eval_loss": eloss.item(),
        #         "perplexity": perplex,
        #         "number of iterations": step + 1,
        #         "Train time per epoch": e_epoch - s_epoch,
        #         "Eval time per epoch": eval_e_epoch - eval_s_epoch,
        #         "avg fwd time": avg_fwd_time / (step + 1),
        #         "avg bwd time": avg_bwd_time / (step + 1),
        #     },
        #     # checkpoint = air.Checkpoint.from_directory(ckpt_path_epoch),
        # )
        


def parse_args():
    
    parser = argparse.ArgumentParser(description="Simple example of training script.")
    parser.add_argument(
        "--mx",
        type=str,
        default="bf16",
        choices=["no", "fp16", "bf16", "fp8"],
        help="Whether to use mixed precision. Choose"
        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
        "and an Nvidia Ampere GPU.",
    )
    
    parser.add_argument(
        "--batch-size-per-device", 
        "-bs", 
        type=int, 
        default=16, 
        help="Batch size to use per device."
    )
    
    parser.add_argument(
        "--eval-batch-size-per-device", 
        type=int, 
        default=16, 
        help="Batch size to use per device (For evaluation)."
    )
    
    parser.add_argument(
        "--num-devices",
        "-nd", 
        type=int, 
        default=4, 
        help="Number of devices to use."
    )
    parser.add_argument(
        "--grad_accum", 
        type=int, 
        default=1, 
        help="Gradient accumulation steps."
    )
    parser.add_argument(
        "--train_path", 
        type=str, 
        help="Path to training jsonl file"
    )
    parser.add_argument(
        "--test_path", 
        type=str, 
        help="Path to testing jsonl file"
    )
    parser.add_argument(
        "--special_token_path", 
        type=str, 
        help="Path to token json file"
    )    
    parser.add_argument(
        "--no-grad-ckpt", 
        action="store_true", 
        help="If passed, will not use gradient checkpointing."
    )
    parser.add_argument(
        "--output_dir", 
        type=str, 
        help="Path to output directory."
    )
    parser.add_argument(
        "--model_name", 
        default="meta-llama/Llama-2-7b-chat-hf", 
        type=str
    )
    parser.add_argument(
        "--num-epochs", 
        type=int, 
        default=1, 
        help="Number of epochs to train for."
    )
    parser.add_argument(
        "--lr", 
        type=float, 
        default=5e-6, 
        help="Learning rate to use."
    )
    
    parser.add_argument(
        "--ctx-len", 
        type=int, 
        default=512, 
        help="Learning rate to use."
    )
    
    parser.add_argument(
        "--as-test", 
        action="store_true",
        help="If passed, will run the script in test mode."
    )
    args = parser.parse_args()
    
    return args

def main():
    
    args = parse_args()
    
    if not args.output_dir:
        raise ValueError("--output_dir must be specified")
    
    config = {
        "lr": args.lr, 
        "num_epochs": args.num_epochs, 
        "seed": 42, 
        "batch_size": args.batch_size_per_device, 
        "gradient_accumulation_steps": args.grad_accum,
        "model_name": args.model_name,
        "block_size": args.ctx_len,
        "eval_batch_size": args.eval_batch_size_per_device,
        "as_test": args.as_test,
    }
    
    ray.init(
        runtime_env={
            "env_vars": {
                "RAY_memory_monitor_refresh_ms": "0",
                "HF_HOME": "/mnt/local_storage/.cache/huggingface",
                "TUNE_RESULT_DIR": args.output_dir,
            }
        }
    )
    
    
    train_ds = create_ray_dataset(args.train_path)
    if args.test_path is not None:
        valid_ds = create_ray_dataset(args.test_path)
    else:
        valid_ds = None
    
    # json file
    with open(args.special_token_path, 'r') as json_file:
        special_tokens = json.load(json_file)["tokens"]
    
    trainer = TorchTrainer(
        training_function,
        train_loop_config={
            "config": config, 
            "args": vars(args), 
            "special_tokens": special_tokens
        },
        run_config=air.RunConfig(
            # storage_path=args.output_dir,
            # # For NFS: disables the default syncing to head node (already disabled on 2.6+)
            # sync_config=tune.SyncConfig(syncer=None),
            storage_path="/mnt/shared_storage/llama-ft",
            checkpoint_config=air.CheckpointConfig(
                num_to_keep=1,
                # Enable distributed checkpointing
                _checkpoint_keep_all_ranks=True,
                _checkpoint_upload_from_workers=True,
            ),
        ),
        scaling_config=air.ScalingConfig(
            num_workers=args.num_devices, 
            use_gpu=True,  
            resources_per_worker={"GPU": 1},
        ),
        datasets={
            "train": train_ds,
            "valid": valid_ds,
        }
    )

    results = trainer.fit()
    
    print("results are stored in:")
    print(results.checkpoint.uri)

    
    
if __name__ == "__main__":
    main() 

You can run this with

python finetune_hf_llm.py \
    -bs 16 \
    -nd 16 \
    --model_name ${MODEL_ID} \
    --output_dir /mnt/cluster_storage/finetuning-llama/ \
    --train_path ./data/train.jsonl \
    --test_path ./data/test.jsonl  \
    --special_token_path ./data/tokens.json \

[BramVanroy]

@kouroshHakha thanks for this. Unfortunately, trying this on a 4x 3090 single-node system did not help (still the mismatch error).

[rileyhun]

I can confirm that I'm also encountering this issue as well using 16 x A10 GPUs, deepspeed==0.8.3, and ray==2.6.1

[BramVanroy]

I tried once more with deepspeed 0.10.1 with the default zero3 config as well as the one provided by @kouroshHakha but to no avail :(

[jaywongs]

I tried once more with deepspeed 0.10.1 with the default zero3 config as well as the one provided by @kouroshHakha but to no avail :(

this pr fixed it：
huggingface/accelerate#1901

[mrwyattii]

this pr fixed it： huggingface/accelerate#1901

It appears this PR fixes the problem. Please reopen if you are still seeing issues. Thanks