Sharded DDP training fails with seq2seq models

[sgugger]

Information

Model I am using (Bert, XLNet ...): T5/BART/mBART/Marian

The problem arises when using:

• the official example scripts: (give details below)
• my own modified scripts: (give details below)

The tasks I am working on is:

• an official GLUE/SQUaD task: seq2seq
• my own task or dataset: (give details below)

To reproduce

Steps to reproduce the behavior:

Run

python -m torch.distributed.launch --nproc_per_node=2 examples/seq2seq/finetune_trainer.py \
--model_name_or_path sshleifer/tiny-mbart --output_dir output_dir --adam_eps 1e-06 --data_dir \
~/Downloads/wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 \
--logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 \
--num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size 4 --sortish_sampler \
--src_lang en_XX --task translation --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 \
--n_train 500 --sharded_ddp

will fail with

Traceback (most recent call last):
File "examples/seq2seq/finetune_trainer.py", line 379, in <module>
main()
File "examples/seq2seq/finetune_trainer.py", line 316, in main
model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
File "/home/sgugger/git/transformers/src/transformers/trainer.py", line 821, in train
self.optimizer.step()
File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/optim/lr_scheduler.py", line 67, in wrapper
return wrapped(*args, **kwargs)
File "/home/sgugger/git/fairscale/fairscale/optim/oss.py", line 210, in step
self._broadcast_params()
File "/home/sgugger/git/fairscale/fairscale/optim/oss.py", line 522, in _broadcast_params
if self.should_bucket_param[param]:
KeyError: Parameter containing:
tensor([[-0.0296,  0.0038],
[ 0.0000,  0.0000],
[ 0.0298,  0.0385],
...,
[-0.0161, -0.0024],
[ 0.0022, -0.0576],
[ 0.0053,  0.0256]], device='cuda:1')
0%|   

Using FP16 also fails.

Expected behavior

The script should run to completion.

[stas00]

This is just a brief log of the 2 distinct errors mentioned in OP:

w/ --fp16 the failure is:

  File "./finetune_trainer.py", line 379, in <module>
    main()
  File "./finetune_trainer.py", line 315, in main
    trainer.train(
  File "/mnt/nvme1/code/huggingface/transformers-master/src/transformers/trainer.py", line 818, in train
    self.scaler.step(self.optimizer)
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/torch/cuda/amp/grad_scaler.py", line 330, in step
    assert len(optimizer_state["found_inf_per_device"]) > 0, "No inf checks were recorded for this optimizer."
AssertionError: No inf checks were recorded for this optimizer.

w/o --fp16 the failure is:

  File "./finetune_trainer.py", line 379, in <module>
    main()
  File "./finetune_trainer.py", line 315, in main
    trainer.train(
  File "/mnt/nvme1/code/huggingface/transformers-master/src/transformers/trainer.py", line 821, in train
    self.optimizer.step()
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/torch/optim/lr_scheduler.py", line 65, in wrapper
    return wrapped(*args, **kwargs)
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/torch/optim/optimizer.py", line 89, in wrapper
    return func(*args, **kwargs)
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/fairscale/optim/oss.py", line 210, in step
    self._broadcast_params()
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/fairscale/optim/oss.py", line 522, in _broadcast_params
    if self.should_bucket_param[param]:
KeyError: Parameter containing:
tensor([[ ...]], device='cuda:1')

It's the very first parameter model.shared.weight in the case of mbart for example.

To test with t5 (same errors), run:

export BS=4; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0   python -m torch.distributed.launch --nproc_per_node=2  ./finetune_trainer.py --model_name_or_path patrickvonplaten/t5-tiny-random --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --sortish_sampler  --task translation_en_XX_to_ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 500 --sharded_ddp 

[blefaudeux]

The first problem (fp16) is easily fixed, it means that the doc is not good enough. Torch's grad scaler is not shard aware (the ranks do not have all the gradients with this technique), but you can use this and that should work.

edit for future readers: that prove wrong, ShardedGradScaler already in use

The second issue is new to me, would you mind sharing a bit more of the reproduction steps ?

[msbaines]

We initialize the should_bucket_param dictionary when the OSS optimizer is created. The assumption is that parameters should be frozen at this point. Any chance parameters are modified after the optimizer was created?

[stas00]

Thank you so much @blefaudeux and @msbaines for your follow up.

To reproduce:

# setup 
git clone  https://github.com/huggingface/transformers
cd transformers
cd examples/seq2seq
wget https://cdn-datasets.huggingface.co/translation/wmt_en_ro.tar.gz 
tar -xzvf wmt_en_ro.tar.gz

to reproduce the 2nd failure w/o --fp16:

export BS=4; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0   python -m torch.distributed.launch --nproc_per_node=2  ./finetune_trainer.py --model_name_or_path  sshleifer/tiny-mbart --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --sortish_sampler --src_lang en_XX --task translation --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 500 --sharded_ddp

and then the first one is to just add --fp16

This is a tiny model that is good enough for testing the mechanics, so no good results to be expected. It's also very quick to download and load. To see real results swap sshleifer/tiny-mbart for sshleifer/distill-mbart-en-ro-12-4.

[stas00]

The first problem (fp16) is easily fixed, it means that the doc is not good enough. Torch's grad scaler is not shard aware (the ranks do not have all the gradients with this technique), but you can use this and that should work.

We are using it already:

transformers/src/transformers/trainer.py

Line 315 in dc9f245

self.scaler = ShardedGradScaler() if self.sharded_dpp else torch.cuda.amp.GradScaler()

if I print the object just before it fails in self.scaler.step(self.optimizer), I get:
<fairscale.optim.grad_scaler.ShardedGradScaler object at 0x7ff27034bac0>

transformers/src/transformers/trainer.py

Line 818 in dc9f245

self.scaler.step(self.optimizer)

FWIW, I experience the exact same issue with deepspeed if I leave trainer's --fp16 code - if I remove it and get deepspeed to handle that the failure goes away. So the common denominator is our code.

[blefaudeux]

Thank you so much @blefaudeux and @msbaines for your follow up.

To reproduce:

# setup 
git clone  https://github.com/huggingface/transformers
cd transformers
cd examples/seq2seq
wget https://cdn-datasets.huggingface.co/translation/wmt_en_ro.tar.gz 
tar -xzvf wmt_en_ro.tar.gz

to reproduce the 2nd failure w/o --fp16:

export BS=4; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0   python -m torch.distributed.launch --nproc_per_node=2  ./finetune_trainer.py --model_name_or_path  sshleifer/tiny-mbart --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --sortish_sampler --src_lang en_XX --task translation --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 500 --sharded_ddp

and then the first one is to just add --fp16

This is a tiny model that is good enough for testing the mechanics, so no good results to be expected. It's also very quick to download and load. To see real results swap sshleifer/tiny-mbart for sshleifer/distill-mbart-en-ro-12-4.

perfect, having a look right now, thanks for the repro help !

[blefaudeux]

We initialize the should_bucket_param dictionary when the OSS optimizer is created. The assumption is that parameters should be frozen at this point. Any chance parameters are modified after the optimizer was created?

good point, it's something that I was planning to address someday actually but not sure how urgent that was

Thank you so much @blefaudeux and @msbaines for your follow up.

To reproduce:

# setup 
git clone  https://github.com/huggingface/transformers
cd transformers
cd examples/seq2seq
wget https://cdn-datasets.huggingface.co/translation/wmt_en_ro.tar.gz 
tar -xzvf wmt_en_ro.tar.gz

to reproduce the 2nd failure w/o --fp16:

export BS=4; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0   python -m torch.distributed.launch --nproc_per_node=2  ./finetune_trainer.py --model_name_or_path  sshleifer/tiny-mbart --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --sortish_sampler --src_lang en_XX --task translation --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 500 --sharded_ddp

and then the first one is to just add --fp16

This is a tiny model that is good enough for testing the mechanics, so no good results to be expected. It's also very quick to download and load. To see real results swap sshleifer/tiny-mbart for sshleifer/distill-mbart-en-ro-12-4.

re: fp16, could it be that the FW has not been run within an AMP context (a) and the scaler is not invoked in the backward (b) ? I cannot find it in the trainer. The scaler needs to be invoked when computing the grads, it will check for infs there which could explain why you're seeing this assert. Note that the hugginface codebase is new to me, so could be that this is wrapped somewhere else and I'm completely missing it

[stas00]

I think the questions you're asking about are all in this training_step code:

transformers/src/transformers/trainer.py

Lines 1126 to 1146 in dc9f245

	if self.use_amp:
	with autocast():
	loss = self.compute_loss(model, inputs)
	else:
	loss = self.compute_loss(model, inputs)
	if self.args.n_gpu > 1:
	loss = loss.mean() # mean() to average on multi-gpu parallel training
	if self.args.gradient_accumulation_steps > 1:
	loss = loss / self.args.gradient_accumulation_steps
	if self.use_amp:
	self.scaler.scale(loss).backward()
	elif self.use_apex:
	with amp.scale_loss(loss, self.optimizer) as scaled_loss:
	scaled_loss.backward()
	else:
	loss.backward()
	return loss.detach()

I didn't write it, but from a quick read it appears that it's a yes to all of your suggestions.

self.use_amp = native amp, use_apex is apex - so we are talking native amp here - that is the branches with use_amp = True

I'll step through with debugger to see that it is actually so.

[blefaudeux]

I think the questions you're asking about are all in this training_step code:

transformers/src/transformers/trainer.py

Lines 1126 to 1146 in dc9f245

	if self.use_amp:
	with autocast():
	loss = self.compute_loss(model, inputs)
	else:
	loss = self.compute_loss(model, inputs)
	if self.args.n_gpu > 1:
	loss = loss.mean() # mean() to average on multi-gpu parallel training
	if self.args.gradient_accumulation_steps > 1:
	loss = loss / self.args.gradient_accumulation_steps
	if self.use_amp:
	self.scaler.scale(loss).backward()
	elif self.use_apex:
	with amp.scale_loss(loss, self.optimizer) as scaled_loss:
	scaled_loss.backward()
	else:
	loss.backward()
	return loss.detach()

I didn't write it, but from a quick read it appears that it's a yes to all of your suggestions.

self.use_amp = native amp, use_apex is apex - so we are talking native amp here - that is the branches with use_amp = True

I'll step through with debugger to see that it is actually so.

looks good indeed ! basically I just know that the "found_inf_per_device" are populated in the "unscale_" step, so this key being absent points to this step being skipped somehow.

[blefaudeux]

@stas00 facebookresearch/fairscale#256 Fixes your repro on a single node, it's a side effect though (bucketing is effectively disabled on a single node), multinode + huggingface is still probably broken. It does look like the model somehow changes after construction, still finding my way around your codebase.

[blefaudeux]

re: --fp16: If the machine is clean it does break, my assumption is that the dist.reduce() in the ShardedGradScaler fails somehow in between the processes. Once it dies, one process stays up actually (visible with nvidia-smi), and all the subsequent runs will work fine. So if that helps it looks like it could be an issue with the dist init, or passing the settings around that.

[stas00]

looks good indeed ! basically I just know that the "found_inf_per_device" are populated in the "unscale_" step, so this key being absent points to this step being skipped somehow.

So, OK, I have it setup in the debugger

It runs unscale_ just fine, but doesn't find any inf:

        inv_scale = self._scale.double().reciprocal().float()
        found_inf = torch.full((1,), 0.0, dtype=torch.float32, device=self._scale.device)

        optimizer_state["found_inf_per_device"] = self._unscale_grads_(optimizer, inv_scale, found_inf, False)

• inv_scale = tensor([1.5259e-05], device='cuda:0')
• found_inf = tensor([0.], device='cuda:0')
• optimizer_state["found_inf_per_device"] = {}

rerunning - inside self._unscale_grads_ we get:

per_device_and_dtype_grads = {defaultdict: 1} defaultdict(<function GradScaler._unscale_grads_.<locals>.<lambda> at 0x7f895c2174c0>, {device(type='cuda', index=1): defaultdict(<class 'list'>, {torch.float32: [tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, 
 cuda:1 = {defaultdict: 1} defaultdict(<class 'list'>, {torch.float32: [tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, 
  default_factory = {type} <class 'list'>
  torch.float32 = {list: 92} [tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([nan, nan], 
   00 = {Tensor: 2} tensor([[nan, nan],\n        [nan, nan]], device='cuda:1')
   01 = {Tensor: 2} tensor([[nan, nan],\n        [nan, nan]], device='cuda:1')
   02 = {Tensor: 2} tensor([[nan, nan],\n        [nan, nan]], device='cuda:1')
   03 = {Tensor: 2} tensor([[nan, nan],\n        [nan, nan]], device='cuda:1')
[...]
   89 = {Tensor: 2} tensor([nan, nan], device='cuda:1')
   90 = {Tensor: 2} tensor([nan, nan], device='cuda:1')
   91 = {Tensor: 2} tensor([-12339.3516, -13527.4590], device='cuda:1')

and then this gets called:

            for device, per_dtype_grads in per_device_and_dtype_grads.items():
                for grads in per_dtype_grads.values():
                    torch._amp_foreach_non_finite_check_and_unscale_(grads,
                                                                     per_device_found_inf.get(device),
                                                                     per_device_inv_scale.get(device))

and it doesn't find anything.

I think the debugger also has a race condition and sometimes we get the first process (cuda:0) and other times the 2nd one (cuda:1).

ok I figured out how to switch threads in pycharm debugger, so I can now go back and force between the 2 processes.

I haven't figured out how to do .to() calls in pycharm debugger yet - when those happen as a step it just hangs so I have carefully to skip over those.

Now rerunning it again and hitting cuda:0 per_device_and_dtype_grads is empty.

And yes, I have noticed that with --fp16 only one process fails with this error - the other keeps on running.

So basically we are having this happen on one process but not the other. Somewhere must be a bug not running the same code for both processes.

(this has just changed in pytorch master - now all subprocesses will die nicely w/o leaving zombies - by yours truly ;)

[blefaudeux]

thanks for the backtrace ! so, for one the fact that the grads are not all the same is expected with this method, the grads are sharded across the ranks (ie: partitioned), depending on which parameters each rank will optimize. The ShardedGradScaler should be aware of that, and syncs in between the ranks to make sure that they all get the same knowledge, looks like this fails somehow then. Having a quick look right now

(well done for the zombie process destruction ! now somehow if the zombie process is still here the next run "works")