Addressing Issue #5241: Updating deepspeed/runtime/zero/stage_1_and_2.py #5252
Conversation
A soft exception case in stage_1_and_2.py/complete_grad_norm_calculation_for_cpu_offload() is not properly handled. Creating a soft logger Warning to better inform the users for further debugging.
|
@desire2020, thanks for the PR. However, a common complaint is the high volume of log messages generated by DeepSpeed and so we need to be really careful in introducing new messages. Also, this PR will break symmetry with norm computation for the non-offload case. To help figure out the right solution here, could you please explain the problem that you are fixing? Thanks! |
| logger.info(f"Warning: invalid gradient detected. Please check your model implementation/configuration to improve the numerical stability.") | ||
| total_norm = -1. |
There was a problem hiding this comment.
Will the following work? Based on the non-offload code path.
| logger.info(f"Warning: invalid gradient detected. Please check your model implementation/configuration to improve the numerical stability.") | |
| total_norm = -1. | |
| total_norm = torch.tensor(-1.0, device=self.device, dtype=torch.float) |
There was a problem hiding this comment.
I think the idea here is to warn the user that this is an unexpected behavior and ususally it would cause/is caused by an error in the gradients. I'm fine with the current temporary fix, but I still hope we can probably throw an exception here or print some warning information to let the user know.
There was a problem hiding this comment.
@desire2020, I was making two points
- It is better for the offload and non-offload code paths to have similar behavior and appearance to users. In other words,
total_normshould be tensors in both paths, and same logging behavior. - Invalid gradient norms would later trigger overflow detection for that iteration. And DS and most frameworks already handle that correctly, with appropriate warning messages. Can you please check that there is no subsequent overflow message warning?
What do you think?
There was a problem hiding this comment.
Hi @tjruwase , yes I checked there is no subsequent overflow message warning. The code immediately crashed, after returning -1 here and then when calculating the grad norm elsewhere, PyTorch throws an uncaught exception when we are trying to calculate the norm of an (autocasted) torch.int64. I was using deepspeed directly with the official huggingface lm example code with mixed precision of bf16, didn't change too much of it. For the error message, please refer to Issue #5241. It's the same in my case.
There was a problem hiding this comment.
@desire2020, thanks for the response. I think you have stumbled on a bigger problem in our code. Normally, we check for overflows before [computing gradient norms]. We assume that the conditions for an overflow are the same as for gradient norm of -1, and so norm computation is skipped on overflows. For bf16 training, we don't check for overflows because we assume overflows are impossible. But it seems your test case contradicts our assumption. Can you try enabling overflow checks for bf16 training to see if overflow is detected? Thanks!
There was a problem hiding this comment.
Hi @tjruwase , sorry if I'm asking a stupid question, is this a currently supported feature or we'd be expecting it in the next release of deepspeed? 'Cause I don't find related implementation of this feature in the current version.
There was a problem hiding this comment.
@desire2020, no worries, I was not clear. I was asking if you could add this feature to your PR? Thanks!
There was a problem hiding this comment.
More specifically, the desired feature is to add "check_overflow" option into the "bf16" dict of the ds_config.
There was a problem hiding this comment.
@tjruwase Ok, do you think I should add it for fp16 at the same time? If people are certain that their objective is stable and simple, they can use this option to further boost their fp16 mixed precision training too.
There was a problem hiding this comment.
Sure, but let's make the fp16 option true by default to preserve BC. Thanks!
@tjruwase - this issue has more information I believe |
|
@desire2020, are you able to complete this PR? Thanks! |
Hi @tjruwase tjruwase, sorry for the late reply. I was busy on some other works. Yes, I will spend some time to work on it asap, and since a new major release is already there, I'm considering starting a new PR to avoid your major changes to config parsing being affected. What do you think? Let me know of your thoughts. |
Either option of a new PR or updating the current is fine with me. Please do whatever is most convenient for you. Thanks for your contribution. |
A soft exception case in stage_1_and_2.py/complete_grad_norm_calculation_for_cpu_offload() is not properly handled. Creating a soft logger Warning to better inform the users for further debugging.
Fixes: #5241