Skip to content

Set environment variables for tpu7x #9586

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

Merged
merged 6 commits into from
Aug 26, 2025
Merged

Set environment variables for tpu7x #9586

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
66dce10
Modification for tpu7x
bhavya01 May 16, 2025
d953d99
Updates for tpu7x
bhavya01 Jun 5, 2025
875811d
Merge branch 'master' into tpu7x
bhavya01 Aug 25, 2025
192d2c2
Update for tpu7x
bhavya01 Aug 25, 2025
d3cd147
Merge branch 'master' into tpu7x
bhavya01 Aug 26, 2025
2a05c40
Fix lint
bhavya01 Aug 26, 2025
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
17 changes: 13 additions & 4 deletions torch_xla/_internal/tpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,8 @@
# Testing only
'0x0056',
'0x0062',
# TPU 7x
'0x0076'
]


Expand Down Expand Up @@ -188,7 +190,10 @@ def version() -> int:
except requests.HTTPError as e:
raise EnvironmentError('Failed to get TPU metadata') from e

match = re.match(r'^v(\d)([A-Za-z]?){7}-(\d+)$', env[xenv.ACCELERATOR_TYPE])
match = re.match(r'^(?:v|tpu)(\d)([A-Za-z]?){7}-(\d+)$',
env[xenv.ACCELERATOR_TYPE])
if not match:
raise EnvironmentError('Failed to parse TPU version from metadata')
return int(match.groups()[0])


Expand Down Expand Up @@ -254,7 +259,8 @@ def configure_topology(local_rank: int,
tpu_env = get_tpu_env()

accelerator_type = tpu_env[xenv.ACCELERATOR_TYPE]
if version() >= 4:
tpu_version = version()
if tpu_version >= 4:
# Process bounds with 4 chips per process
default_process_bounds = MeshShape.from_string(
tpu_env[xenv.TPU_PROCESS_BOUNDS])
Expand All @@ -270,8 +276,11 @@ def configure_topology(local_rank: int,
process_bounds = default_process_bounds * chips_per_process

os.environ.setdefault(xenv.TPU_CHIPS_PER_PROCESS_BOUNDS, '1,1,1')
os.environ.setdefault(xenv.TPU_PROCESS_BOUNDS,
','.join(str(dim) for dim in process_bounds))
process_bounds_str = ','.join(str(dim) for dim in process_bounds)
if tpu_version == 7:
process_bounds_str += ',2'

os.environ.setdefault(xenv.TPU_PROCESS_BOUNDS, process_bounds_str)

# Assume each TPU has the same number of local processes with the same ports
worker_id = int(tpu_env[xenv.WORKER_ID])
Expand Down