Use the val argument in mask_nan_or_inf_with_val_inplace

[vineethsaivs]

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.

[ebarkhordar]

get_grad_norm_direct in zero/stage3.py never got folded into the helper, so the three copies in the description are missing one. It still masks inline, with the same hardcoded literal:

# deepspeed/runtime/zero/stage3.py:2311-2316
norm_is_inf = total_norm.isinf()
norm_is_nan = total_norm.isnan()
inf_or_nan = norm_is_nan.logical_or(norm_is_inf)

err = torch.tensor(-1.0, device=self.device, dtype=torch.float)
total_norm = torch.where(inf_or_nan, err, total_norm)

So that path cannot honor val either. Swapping those six lines for mask_nan_or_inf_with_val_inplace(total_norm, device=self.device) returns the same value at the same dtype: total_norm comes fresh off total_norm_cuda ** (1. / norm_type), so the in-place write has no other reader, and I ran torch.where against masked_fill_ for both get_norm_dtype() cases, fp64 and fp32, and got -1.0 at the input dtype from each.

Your three call sites are right. I enumerated them with ast rather than grep and got exactly runtime/utils.py:963, zero/stage3.py:1885 and zero/stage_1_and_2.py:2114, all taking the default.

Not blocking, and it may well be cleaner as its own PR.

[vineethsaivs]

You are right, that is a fourth site and it cannot honour val either. Two things to add before it moves.

The "no other reader" argument needs the other branch too. get_grad_norm_direct reaches that masking from two places, and only one of them produces a fresh tensor:

# norm_type == inf
total_norm = total_norm_cuda[0]            # a view into total_norm_cuda
# else
total_norm = total_norm_cuda**(1. / norm_type)   # fresh

So masked_fill_ on the inf path also writes into total_norm_cuda. It is still safe, because nothing reads total_norm_cuda after that line and the function returns immediately, but the justification should say so rather than resting on the finite-norm path alone.

Dtypes agree on both paths. The inf branch comes off FloatTensor, so fp32; the finite branch off get_norm_dtype(), fp32 or fp64. masked_fill_ keeps the input dtype and torch.where promotes to the same thing, which matches what you measured.

Taking your suggestion: this PR stays the one-line val fix. Happy to fold the stage3 site in here instead if you would rather have one PR, just say which you prefer.

[ebarkhordar]

Separate, I'd say. This one is a one line change with a self-contained argument. The stage3 swap turns an allocating torch.where into an in place write on total_norm_cuda[0], which is the aliasing you just described, so it wants its own diff rather than riding along under this title.

You are right that my justification only covered the else branch. total_norm = total_norm_cuda[0] is a view, so on the inf path the write lands in total_norm_cuda; nothing reads it afterwards and the function returns on the next line, but that is the argument there, not the fresh-tensor one.

[vineethsaivs]

Agreed on both counts. Keeping this one at the single line, and I will open the stage3 torch.where change separately so the aliasing argument gets its own diff rather than hiding under this title.