Skip to content

Use the val argument in mask_nan_or_inf_with_val_inplace - #8386

Open
vineethsaivs wants to merge 1 commit into
deepspeedai:masterfrom
vineethsaivs:fix/mask-nan-inf-val
Open

Use the val argument in mask_nan_or_inf_with_val_inplace#8386
vineethsaivs wants to merge 1 commit into
deepspeedai:masterfrom
vineethsaivs:fix/mask-nan-inf-val

Conversation

@vineethsaivs

Copy link
Copy Markdown
Contributor

What

mask_nan_or_inf_with_val_inplace(input, device=None, val=-1.) accepts val and then builds the fill tensor from a hardcoded -1.0, so the argument is dead:

t = torch.tensor([float('nan'), 2.0])
mask_nan_or_inf_with_val_inplace(t, device=t.device, val=0.0)
# t -> tensor([-1.,  2.])   expected tensor([0., 2.])

Why it is there

val arrived with the helper in #7184, which folded three copies of the inf/nan masking (runtime/utils.py, zero/stage3.py, zero/stage_1_and_2.py) into one function. The parameter is the knob that refactor introduced, and the body kept the literal from the code it replaced.

Fix

One line: build err from val.

All three in-tree callers use the default, and -1. is the same float the body hardcoded, so the gradient-norm paths are bit-identical.

Test

New test_mask_nan_or_inf_with_val_inplace_honors_val in tests/unit/runtime/test_runtime_utils.py, covering nan / +inf / -inf for several vals plus the unchanged default.

$ python -m pytest unit/runtime/test_runtime_utils.py -k "mask_nan_or_inf or call_to_str or count_used_parameters"
# before: 1 failed, 2 passed
#   E  assert [-1.0, -1.0, -1.0, 2.0] == [0.0, 0.0, 0.0, 2.0]
# after:  3 passed

The TestClipGradNorm* / TestCheckOverflow cases in the same file are DistributedTest and were deselected on this CPU-only box; they do not touch val.

The helper takes `val` and then hardcodes -1.0 in the tensor it fills with,
so any caller asking for a different sentinel silently gets -1 instead.

`val` was added with the helper in deepspeedai#7184, which folded three copies of the
inf/nan masking into one function; the parameter is the knob that refactor
introduced and it was never wired to the body. Every in-tree caller uses the
default, so the fix is a no-op for them: -1. is the same float the body
hardcoded.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant