deepspeedai / deepspeedai/DeepSpeed
[BUG] Mismatch between dtype settings in model and ds_config results in NaN loss
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 43.1k
- Forks
- 5k
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 112
Description
Describe the bug
When there is a mismatch between the dtype settings of the model and ds_config, training starts without any specific error and the loss turns NaN (this issue occurs mainly in Zero stage0).
I suggest adding a dtype check between the model and config during the execution of deepspeed.initialize and throwing an assert if they do not match. What do you think?
To Reproduce
- Use the DeepSpeedExample with cifar.
- Edit cifar10_deepspeed.py as follows:
+ net = net.half()
model_engine, optimizer, trainloader, __ = deepspeed.initialize(
args=args,
model=net,
model_parameters=parameters,
training_data=trainset,
config=ds_config,
)
# Get the local device name (str) and local rank (int).
local_device = get_accelerator().device_name(model_engine.local_rank)
local_rank = model_engine.local_rank
# For float32, target_dtype will be None so no datatype conversion needed.
target_dtype = None
if model_engine.bfloat16_enabled():
target_dtype = torch.bfloat16
elif model_engine.fp16_enabled():
target_dtype = torch.half
+ target_dtype = torch.half
- Execute the following:
$ deepspeed --bind_cores_to_rank cifar10_deepspeed.py --dtype fp32 --stage 0
- You will observe that the loss turns NaN.
[ 1, 2000] loss: nan
[ 2, 2000] loss: nan
[ 3, 2000] loss: nan
[ 4, 2000] loss: nan
[ 5, 2000] loss: nan
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with deepspeed.initialize and reproduce the issue using cifar10_deepspeed.py, ds_config, and the shown deepspeed command. Trace how the model dtype and configuration dtype are handled, especially for Zero stage 0. Done means a mismatch is detected during initialization instead of allowing training to produce NaN loss.
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
- Mostly clear
- Newbie friendliness
- 45/100