Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix XLA fp16 and bf16 error checking #18913

Merged
merged 2 commits into from
Sep 7, 2022
Merged

Fix XLA fp16 and bf16 error checking #18913

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5865dfe
Fix XLA fp16 and bf16 error checking
ymwangg Sep 7, 2022
5be4119
Update src/transformers/training_args.py
ymwangg Sep 7, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions src/transformers/training_args.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,15 @@ def get_int_from_env(env_keys, default):
return default


def get_xla_device_type(device: "torch.device") -> Optional[str]:
"""
Returns the xla device type (CPU|GPU|TPU) or None if the device is a non-xla device.
"""
if is_torch_tpu_available():
return xm.xla_real_devices([device])[0].split(":")[0]
return None


class OptimizerNames(ExplicitEnum):
"""
Stores the acceptable string identifiers for optimizers.
Expand Down Expand Up @@ -1108,7 +1117,7 @@ def __post_init__(self):
self.framework == "pt"
and is_torch_available()
and (self.device.type != "cuda")
and not (self.device.type == "xla" and "GPU_NUM_DEVICES" in os.environ)
and (get_xla_device_type(self.device) != "GPU")
and (self.fp16 or self.fp16_full_eval)
):
raise ValueError(
Expand All @@ -1120,7 +1129,7 @@ def __post_init__(self):
self.framework == "pt"
and is_torch_available()
and (self.device.type != "cuda")
and not (self.device.type == "xla" and "GPU_NUM_DEVICES" in os.environ)
and (get_xla_device_type(self.device) != "GPU")
and (self.device.type != "cpu")
and (self.bf16 or self.bf16_full_eval)
):
Expand Down