Add 0dim Tensor overload for _foreach_div

[janeyx99]

This PR is ALMOST basically just following the steps from #106677 EXCEPT! We do add one feature. Similar to fused_adam(w), for the CUDA dispatches: when the scalar tensor is on CPU, we .item and redispatch to the normal scalar overload. Otherwise, the cuda kernel will complain about mismatch in devices between the scalar and the tensors.

Why do we add this feature? Our optimizers want to allow lr as a tensor, and lr could be a CPU tensor. lr is used with foreach_div_ in Adam, so our CI will break otherwise.

After this PR, _foreach_mul and _foreach_div will accept either a CPU or a GPU tensor for the scalar tensor (vs only a GPU tensor). They join the ranks of fused_adam(w) in this characteristic. I did not yet do the same thing for foreach_add (the only other foreach op with a .Tensor overload) because there is no use case and will be more involved.

cc @crcrpar

Stack from ghstack (oldest at bottom):

• -> Add 0dim Tensor overload for _foreach_div #113688

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/113688

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit ed39357 with merge base afef32b ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[janeyx99]

rename to expect_fastpath, which is what it really is and is less confusing

[janeyx99]

This works now, but made me realize I didn't add a case for _foreach_add when I added the overload. Adding that now.

[albanD]

Please confirm there is a test case that also ensures that a list of cpu Tensor and a list of cuda Tensor properly raises.

[janeyx99]

Please confirm there is a test case that also ensures that a list of cpu Tensor and a list of cuda Tensor properly raises.

test_parity tests that both a forloop over the ops + the foreach op will return the same thing or error the same:

pytorch/test/test_foreach.py

Lines 163 to 175 in 42b2b9e

	try:
	with ctxmgr:
	actual = func([sample.input, *sample.args], self.is_cuda, expect_fastpath, **kwargs)
	except Exception as e:
	with (
	self.assertRaisesRegex(type(e), re.escape(str(e)))
	if not (op.has_no_in_place or op.has_no_out_of_place)
	else self.assertRaises(type(e))
	):
	ref([ref_input, *sample.ref_args], **ref_kwargs)
	else:
	expected = ref([ref_input, *sample.ref_args], **ref_kwargs)
	self.assertEqual(expected, actual)

Note that if the reference code errors in line 174, the test will fail.

[janeyx99]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[mlazos]

Please confirm there is a test case that also ensures that a list of cpu Tensor and a list of cuda Tensor properly raises.

@janeyx99 hmm at what level will this error, I probably need to add the grouping at the dynamo level if this occurs when FakeTensor tracing as we discussed

[janeyx99]

@janeyx99 hmm at what level will this error, I probably need to add the grouping at the dynamo level if this occurs when FakeTensor tracing as we discussed

Ah, what happens is that we'll dispatch into the CUDA impl foreach_tensor_div_tensor_kernel_cuda which is a C++ function that checks device types. Since the device types do not match, we then dispatch to the slow impl, which is just another C++ function that for loops over the normal div. Once it calls into div, it will error. It's worth ensuring what torch.compile will do (hopefully it's not too smart; we don't want it converting the second list to GPU tensors).

[mlazos]

@janeyx99 hmm at what level will this error, I probably need to add the grouping at the dynamo level if this occurs when FakeTensor tracing as we discussed

Ah, what happens is that we'll dispatch into the CUDA impl foreach_tensor_div_tensor_kernel_cuda which is a C++ function that checks device types. Since the device types do not match, we then dispatch to the slow impl, which is just another C++ function that for loops over the normal div. Once it calls into div, it will error. It's worth ensuring what torch.compile will do (hopefully it's not too smart; we don't want it converting the second list to GPU tensors).

Nah this is good, this is exactly what I want. I also will only apply the special handling we discussed when there is a single tensor in the second arg.