Skip to content

Fix norm data-type when using AMP.#7878

Merged
JackCaoG merged 2 commits intomasterfrom
ysiraichi/fix-norm-amp-dtype
Aug 19, 2024
Merged

Fix norm data-type when using AMP.#7878
JackCaoG merged 2 commits intomasterfrom
ysiraichi/fix-norm-amp-dtype

Conversation

@ysiraichi
Copy link
Copy Markdown
Collaborator

This PR fixes the result of norm operation when using AMP.

The cast policy defined in autocast_mode.cpp for norm.ScalarOpt_dim is: fp32_append_dtype. Which means that it will forward the call to norm.ScalarOpt_dim_dtype, appending at::kFloat (i.e. float32) as its last argument. Such an argument represents the data-type of the result tensor.

Even though we were correctly lowering the operation so as to return a float32 tensor, upon checking, the result tensor actually inherited its data-type from the input. The solution was to call XLATensor::CreateFrom with at::kFloat argument. The example below illustrates the problem:

>>> x = torch.rand((10, 10), dtype=torch.float16, device="xla")
>>> with torch.cuda.amp.autocast(dtype=torch.float16):
        r = torch.norm(x, p=2, dim=1)

>>> r
# HLO representation shows it's returning a f32 tensor, though.
tensor(..., dtype=torch.float16)

cc @miladm @JackCaoG

@JackCaoG JackCaoG merged commit ac13bf2 into master Aug 19, 2024
@JackCaoG JackCaoG deleted the ysiraichi/fix-norm-amp-dtype branch August 19, 2024 18:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants