# Incorrect total train batch size when using tp\_size \> 1 and deepspeed

**URL:** <https://discuss.huggingface.co/t/incorrect-total-train-batch-size-when-using-tp-size-1-and-deepspeed/156045>\
**Category:** DeepSpeed\
**Created:** [May 20, 2025, 2:53pm UTC](https://discuss.huggingface.co/t/incorrect-total-train-batch-size-when-using-tp-size-1-and-deepspeed/156045 "2025-05-20T14:53:07Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![cynricfu](https://avatars.discourse-cdn.com/v4/letter/c/4da419/32.png) [@cynricfu](https://discuss.huggingface.co/u/cynricfu)\
**Post date:** [May 20, 2025, 2:53pm UTC](https://discuss.huggingface.co/t/incorrect-total-train-batch-size-when-using-tp-size-1-and-deepspeed/156045/1 "2025-05-20T14:53:08Z")

</div>

When I launch a job to run [run\_clm.py](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_clm.py) on two 8-GPU nodes with `tp_size = 2`, `per_device_train_batch_size = 1` (set via TrainingArguments) and deepspeed zero2 (actually, it does not seem to matter which zero stage I choose), I got the following log which does not seem correct.

```auto
node032: [INFO|trainer.py:2414] 2025-05-20 22:25:24,210 >> *****Running training*****
node032: [INFO|trainer.py:2415] 2025-05-20 22:25:24,210 >> Num examples = 287
node032: [INFO|trainer.py:2416] 2025-05-20 22:25:24,210 >> Num Epochs = 3
node032: [INFO|trainer.py:2417] 2025-05-20 22:25:24,210 >> Instantaneous batch size per device = 1
node032: [INFO|trainer.py:2420] 2025-05-20 22:25:24,210 >> Total train batch size (w. parallel, distributed & accumulation) = 16
node032: [INFO|trainer.py:2421] 2025-05-20 22:25:24,210 >> Gradient Accumulation steps = 1
node032: [INFO|trainer.py:2422] 2025-05-20 22:25:24,210 >> Total optimization steps = 108
node032: [INFO|trainer.py:2423] 2025-05-20 22:25:24,211 >> Number of trainable parameters = 3,085,938,688
node028: [INFO|trainer.py:2414] 2025-05-20 22:25:27,491 >> *****Running training*****
node028: [INFO|trainer.py:2415] 2025-05-20 22:25:27,491 >> Num examples = 287
node028: [INFO|trainer.py:2416] 2025-05-20 22:25:27,491 >> Num Epochs = 3
node028: [INFO|trainer.py:2417] 2025-05-20 22:25:27,491 >> Instantaneous batch size per device = 1
node028: [INFO|trainer.py:2420] 2025-05-20 22:25:27,491 >> Total train batch size (w. parallel, distributed & accumulation) = 16
node028: [INFO|trainer.py:2421] 2025-05-20 22:25:27,491 >> Gradient Accumulation steps = 1
node028: [INFO|trainer.py:2422] 2025-05-20 22:25:27,491 >> Total optimization steps = 108
node028: [INFO|trainer.py:2423] 2025-05-20 22:25:27,492 >> Number of trainable parameters = 3,085,938,688

```

Based on my understanding, with 16 GPUs in total and `tp_size = 2` and `per_device_train_batch_size = 1`, the total train batch size should be 8 instead of 16, which is quite weird to me.

The training can successfully run without error though.

How does this happen? Is there anything that I could do wrong?

My environment:

```auto
transformers==4.51.3
accelerate==1.6.0
deepspeed==0.16.7
torch==2.6.0+cu124

```

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [May 20, 2025, 3:10pm UTC](https://discuss.huggingface.co/t/incorrect-total-train-batch-size-when-using-tp-size-1-and-deepspeed/156045/2 "2025-05-20T15:10:31Z")

</div>

Different but similar issue.

> <https://github.com/huggingface/accelerate/issues/2712>
>
> \### System Info
> 
> \`\`\`Shell
> \- \`Accelerate\` version: 0.29.3
> \- Platform: Linux-5….15.0-104.119.4.2.el8uek.x86\_64-x86\_64-with-glibc2.31
> \- \`accelerate\` bash location: /tmp/ray/session\_2024-04-25\_09-37-01\_001890\_8/runtime\_resources/pip/9ef89d7a93ad0f6658e2ddb0a14f262115166bb8/virtualenv/bin/accelerate
> \- Python version: 3.9.19
> \- Numpy version: 1.24.4
> \- PyTorch version (GPU?): 2.0.1+cu118 (True)
> \- PyTorch XPU available: False
> \- PyTorch NPU available: False
> \- PyTorch MLU available: False
> \- System RAM: 2015.21 GB
> \- GPU type: NVIDIA A100-SXM4-40GB
> \- \`Accelerate\` default config:
> Not found
> \`\`\`
> 
> 
> \### Information
> 
> \- \[\] The official example scripts
> \- \[X\] My own modified scripts
> 
> \### Tasks
> 
> \- \[\] One of the scripts in the examples/ folder of Accelerate or an officially supported \`no\_trainer\` script in the \`examples\` folder of the \`transformers\` repo (such as \`run\_no\_trainer\_glue.py\`)
> \- \[X\] My own task or dataset (give details below)
> 
> \### Reproduction
> 
> I am using Ray Train with \`TorchTrainer\` on a Ray Cluster, DeepSpeed, and the HF transformers \`Trainer\`, which calls accelerate implicitly.
> 
> \`requirements.txt\`:
> \`\`\`
> transformers\[deepspeed\]==4.40.1
> accelerate==0.29.3
> deepspeed==0.14.2
> datasets
> flash-attn
> ray\[train\]==2.11.0
> \`\`\`
> 
> \`train.py\`
> \`\`\`py
> from dataclasses import dataclass, field
> import os
> from datetime import datetime
> 
> import torch
> from transformers import HfArgumentParser, TrainingArguments, Trainer, AutoModelForCausalLM, AutoTokenizer, DataCollatorForLanguageModeling
> from datasets import load\_dataset
> 
> from ray.train import ScalingConfig, RunConfig
> from ray.train.torch import TorchTrainer
> from ray.train.huggingface.transformers import RayTrainReportCallback, prepare\_trainer
> 
> @dataclass
> class ModelArguments:
> model: str = field(default="meta-llama/Llama-2-7b-chat-hf", metadata={"help": "Model path or HF hub ID."})
> use\_flash\_attn: bool = field(default=True, metadata={"help": "Enable Flash Attention in training."})
> 
> 
> @dataclass
> class DataArguments:
> dataset: str = field(metadata={"help": "Dataset path or HF hub ID."})
> 
> 
> @dataclass
> class RayArguments:
> num\_workers: int = field(default=1, metadata={"help": "Number of workers to use in Ray cluster."})
> 
> 
> def train(training\_args: TrainingArguments, model\_args: ModelArguments, data\_args: DataArguments):
> # bug when initializing distributed state with ray train + transformers
> # https://github.com/ray-project/ray/issues/44204
> # need to manually call deepspeed.init\_distributed() here to get the correct world\_size
> import deepspeed
> deepspeed.init\_distributed()
> 
> # initialize model, tokenizer
> model = AutoModelForCausalLM.from\_pretrained(
> model\_args.model,
> trust\_remote\_code=True,
> token=True,
> attn\_implementation="flash\_attention\_2" if model\_args.use\_flash\_attn else "eager",
> torch\_dtype=torch.bfloat16 if model\_args.use\_flash\_attn else torch.float32,
> )
> tokenizer = AutoTokenizer.from\_pretrained(model\_args.model, trust\_remote\_code=True)
> if tokenizer.pad\_token is None:
> tokenizer.pad\_token = tokenizer.eos\_token
> 
> # preprocess dataset
> dataset = load\_dataset(data\_args.dataset)
> dataset = dataset.map(
> lambda examples: tokenizer(examples\["text"\]),
> batched=True,
> remove\_columns=dataset\["train"\].column\_names
> )
> train\_dataset = dataset\["train"\]
> val\_dataset = dataset\["validation"\]
> data\_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
> 
> # prepare trainer
> trainer = Trainer(
> model=model,
> tokenizer=tokenizer,
> args=training\_args,
> train\_dataset=train\_dataset,
> eval\_dataset=val\_dataset,
> data\_collator=data\_collator,
> )
> trainer.add\_callback(RayTrainReportCallback())
> trainer = prepare\_trainer(trainer)
> 
> # run training
> trainer.train()
> 
> 
> if \_\_name\_\_ == "\_\_main\_\_":
> parser = HfArgumentParser((TrainingArguments, ModelArguments, DataArguments, RayArguments))
> training\_args, model\_args, data\_args, ray\_args = parser.parse\_args\_into\_dataclasses()
> 
> scaling\_config = ScalingConfig(
> num\_workers=ray\_args.num\_workers,
> use\_gpu=True,
> )
> timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
> run\_config = RunConfig(
> storage\_path=training\_args.output\_dir,
> name=f"train-{timestamp}",
> )
> trainer = TorchTrainer(
> lambda: train(training\_args, model\_args, data\_args),
> scaling\_config=scaling\_config,
> run\_config=run\_config,
> )
> result = trainer.fit()
> print(result)
> print(result.best\_checkpoints)
> \`\`\`
> \`runtime-env.json\`
> \`\`\`json
> {
> "pip": "requirements.txt"
> }
> \`\`\`
> \`deepspeed-zero-3.json\`
> \`\`\`json
> {
> "train\_batch\_size": "auto",
> "train\_micro\_batch\_size\_per\_gpu": "auto",
> "gradient\_accumulation\_steps": 1,
> "zero\_optimization": {
> "stage": 3,
> "offload\_optimizer": {
> "device": "none"
> },
> "offload\_param": {
> "device": "none"
> },
> "stage3\_gather\_16bit\_weights\_on\_model\_save": true
> },
> "gradient\_clipping": 1.0,
> "steps\_per\_print": "inf",
> "bf16": {
> "enabled": true
> }
> }
> \`\`\`
> \`run-train.sh\`
> \`\`\`bash
> ray job submit \\
> --working-dir . \\
> --runtime-env runtime-env.json \\
> -- \\
> python3 train.py \\
> --model "EleutherAI/gpt-neo-2.7B" \\
> --dataset "OpenAssistant/oasst2" \\
> --bf16 True \\
> --tf32 True \\
> --output\_dir /mnt/output \\
> --num\_train\_epochs 2 \\
> --per\_device\_train\_batch\_size 2 \\
> --per\_device\_eval\_batch\_size 2 \\
> --gradient\_accumulation\_steps 1 \\
> --evaluation\_strategy "steps" \\
> --eval\_steps 20 \\
> --save\_strategy "steps" \\
> --save\_steps 20 \\
> --save\_total\_limit 3 \\
> --logging\_steps 1 \\
> --learning\_rate 2e-5 \\
> --warmup\_ratio 0.03 \\
> --report\_to tensorboard \\
> --deepspeed deepspeed-zero-3.json \\
> --use\_flash\_attn True \\
> --num\_workers 2
> \`\`\`
> 
> I get the following error message:
> \`\`\`
> ray.exceptions.RayTaskError(AssertionError): ray::\_RayTrainWorker\_\_execute.get\_next() (pid=85669, ip=10.244.14.216, actor\_id=74d3ccf4a9a71a9c05c6e3ed13000000, repr=\<ray.train.\_internal.worker\_group.RayTrainWorker object at 0x7fee3805dd30\>)
> File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/train/\_internal/worker\_group.py", line 33, in \_\_execute
> raise skipped from exception\_cause(skipped)
> File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/train/\_internal/utils.py", line 169, in discard\_return\_wrapper
> train\_func(\*args, \*\*kwargs)
> File "/tmp/ray/session\_2024-04-25\_09-37-01\_001890\_8/runtime\_resources/working\_dir\_files/\_ray\_pkg\_071f1221244a3868/soundhound/polarisnlu/llm/train.py", line 110, in \<lambda\>
> lambda: train(training\_args, model\_args, data\_args),
> File "/tmp/ray/session\_2024-04-25\_09-37-01\_001890\_8/runtime\_resources/working\_dir\_files/\_ray\_pkg\_071f1221244a3868/soundhound/polarisnlu/llm/train.py", line 91, in train
> trainer.train()
> File "/tmp/ray/session\_2024-04-25\_09-37-01\_001890\_8/runtime\_resources/pip/9ef89d7a93ad0f6658e2ddb0a14f262115166bb8/virtualenv/lib/python3.9/site-packages/transformers/trainer.py", line 1859, in train
> return inner\_training\_loop(
> File "/tmp/ray/session\_2024-04-25\_09-37-01\_001890\_8/runtime\_resources/pip/9ef89d7a93ad0f6658e2ddb0a14f262115166bb8/virtualenv/lib/python3.9/site-packages/transformers/trainer.py", line 2012, in \_inner\_training\_loop
> model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
> File "/tmp/ray/session\_2024-04-25\_09-37-01\_001890\_8/runtime\_resources/pip/9ef89d7a93ad0f6658e2ddb0a14f262115166bb8/virtualenv/lib/python3.9/site-packages/accelerate/accelerator.py", line 1266, in prepare
> result = self.\_prepare\_deepspeed(\*args)
> File "/tmp/ray/session\_2024-04-25\_09-37-01\_001890\_8/runtime\_resources/pip/9ef89d7a93ad0f6658e2ddb0a14f262115166bb8/virtualenv/lib/python3.9/site-packages/accelerate/accelerator.py", line 1652, in \_prepare\_deepspeed
> engine, optimizer, \_, lr\_scheduler = deepspeed.initialize(\*\*kwargs)
> File "/tmp/ray/session\_2024-04-25\_09-37-01\_001890\_8/runtime\_resources/pip/9ef89d7a93ad0f6658e2ddb0a14f262115166bb8/virtualenv/lib/python3.9/site-packages/deepspeed/\_\_init\_\_.py", line 167, in initialize
> config\_class = DeepSpeedConfig(config, mpu)
> File "/tmp/ray/session\_2024-04-25\_09-37-01\_001890\_8/runtime\_resources/pip/9ef89d7a93ad0f6658e2ddb0a14f262115166bb8/virtualenv/lib/python3.9/site-packages/deepspeed/runtime/config.py", line 795, in \_\_init\_\_
> self.\_configure\_train\_batch\_size()
> File "/tmp/ray/session\_2024-04-25\_09-37-01\_001890\_8/runtime\_resources/pip/9ef89d7a93ad0f6658e2ddb0a14f262115166bb8/virtualenv/lib/python3.9/site-packages/deepspeed/runtime/config.py", line 978, in \_configure\_train\_batch\_size
> self.\_batch\_assertion()
> File "/tmp/ray/session\_2024-04-25\_09-37-01\_001890\_8/runtime\_resources/pip/9ef89d7a93ad0f6658e2ddb0a14f262115166bb8/virtualenv/lib/python3.9/site-packages/deepspeed/runtime/config.py", line 926, in \_batch\_assertion
> assert train\_batch == micro\_batch \* grad\_acc \* self.world\_size, (
> AssertionError: Check batch related parameters. train\_batch\_size is not equal to micro\_batch\_per\_gpu \* gradient\_acc\_step \* world\_size 2 != 2 \* 1 \* 2
> \`\`\`
> 
> It says \`train\_batch\_size\` is set incorrectly, even though I didn't set it manually.
> 
> \### Expected behavior
> 
> \`train\_batch\_size\` should be set to 4 in this case, since my per-GPU batch size is 2, and I have 2 GPUs.
> 
> I did a bit of digging and I think \`train\_batch\_size\` is incorrectly set in \`Accelerator.\_prepare\_deepspeed\`:
> https://github.com/huggingface/accelerate/blob/6af157ea93dfbace1db88b0fdc7dfb568dfdd5a5/src/accelerate/accelerator.py#L1537-1539
> 
> \`\`\`py
> config\_kwargs = {
> "train\_micro\_batch\_size\_per\_gpu": batch\_size\_per\_device,
> "train\_batch\_size": batch\_size\_per\_device
> \`\`\`
> 
> It shouldn't be set to the same value as \`train\_micro\_batch\_size\_per\_gpu\`, since if there are multiple GPUs then this would be incorrect.
> 
> I was able to fix the problem by resetting \`train\_batch\_size\` to None with a very hacky code injection fix:
> \`\`\`py
> def train():
> # ...
> import functools
> def fix\_train\_batch\_size(function):
> @functools.wraps(function)
> def wrapped\_function(\*args, config\_params=None, \*\*kwargs):
> if isinstance(config\_params, dict):
> config\_params\['train\_batch\_size'\] = None
> print('Reset train\_batch\_size in config\_params')
> return function(\*args, config\_params=config\_params, \*\*kwargs)
> return wrapped\_function
> deepspeed.initialize = fix\_train\_batch\_size(deepspeed.initialize)
> \`\`\`

> **[Question about the calculation of \`train\_batch\_size\` · deepspeedai DeepSpeed...](https://github.com/deepspeedai/DeepSpeed/discussions/3876)**
>
> In the doc https://www.deepspeed.ai/docs/config-json/#batch-size-related-parameters, train\_batch\_size is calculated in this way: train\_batch\_size = train\_micro\_batch\_size\_per\_gpu \* gradient\_accumul...