Skip to content

Catch info from metadata #4724

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 16 commits into from
Mar 28, 2023
Merged

Catch info from metadata #4724

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
2420d78
try version modify first and may try to change to tpu_env['ACCELERATO…
ManfeiBai Mar 7, 2023
571fa91
catch like patches
ManfeiBai Mar 7, 2023
18eba79
add test
ManfeiBai Mar 8, 2023
db330c0
update to patches one
ManfeiBai Mar 8, 2023
2bcbae9
change match to match new
ManfeiBai Mar 8, 2023
4946290
update space
ManfeiBai Mar 8, 2023
25791ef
delete for abcd
ManfeiBai Mar 16, 2023
76ab2b3
update to match more
ManfeiBai Mar 16, 2023
2d7aea2
update code clean
ManfeiBai Mar 16, 2023
0eecb86
str to int
ManfeiBai Mar 16, 2023
eb48e50
catch str
ManfeiBai Mar 16, 2023
92a84e5
updated to match first
ManfeiBai Mar 16, 2023
44c4bbf
match
ManfeiBai Mar 17, 2023
0f1ca51
update to ol;d
ManfeiBai Mar 22, 2023
6d91493
update to ole
ManfeiBai Mar 22, 2023
620bdba
format
ManfeiBai Mar 22, 2023
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
57 changes: 57 additions & 0 deletions test/pjrt/test_experimental_tpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,63 @@ def test_get_worker_ips(self, worker_network_endpoints, expected):
self.assertListEqual(worker_ips, expected)

@parameterized.named_parameters(
('v5-4_process_0', {
'ACCELERATOR_TYPE': 'v5-4',
xenv.TPU_PROCESS_BOUNDS: '2,2,1',
xenv.TPU_CHIPS_PER_PROCESS_BOUNDS: '1,1,1',
'WORKER_ID': '0'
}, ['localhost'], 0, 4, {
xenv.TPU_CHIPS_PER_PROCESS_BOUNDS:
'1,1,1',
xenv.TPU_PROCESS_BOUNDS:
'2,2,1',
xenv.CLOUD_TPU_TASK_ID:
'0',
xenv.TPU_PROCESS_PORT:
'8476',
xenv.TPU_PROCESS_ADDRESSES:
'localhost:8476,localhost:8477,localhost:8478,localhost:8479',
xenv.TPU_VISIBLE_CHIPS:
'0',
}),
('v5abcdefg-4_process_0', {
'ACCELERATOR_TYPE': 'v5abcdefg-4',
xenv.TPU_PROCESS_BOUNDS: '2,2,1',
xenv.TPU_CHIPS_PER_PROCESS_BOUNDS: '1,1,1',
'WORKER_ID': '0'
}, ['localhost'], 0, 4, {
xenv.TPU_CHIPS_PER_PROCESS_BOUNDS:
'1,1,1',
xenv.TPU_PROCESS_BOUNDS:
'2,2,1',
xenv.CLOUD_TPU_TASK_ID:
'0',
xenv.TPU_PROCESS_PORT:
'8476',
xenv.TPU_PROCESS_ADDRESSES:
'localhost:8476,localhost:8477,localhost:8478,localhost:8479',
xenv.TPU_VISIBLE_CHIPS:
'0',
}),
('v5abcdefg-16_process_0', {
'ACCELERATOR_TYPE': 'v5abcdefg-16',
xenv.TPU_PROCESS_BOUNDS: '2,2,1',
xenv.TPU_CHIPS_PER_PROCESS_BOUNDS: '1,1,1',
'WORKER_ID': '0'
}, ['localhost'], 0, 4, {
xenv.TPU_CHIPS_PER_PROCESS_BOUNDS:
'1,1,1',
xenv.TPU_PROCESS_BOUNDS:
'2,2,1',
xenv.CLOUD_TPU_TASK_ID:
'0',
xenv.TPU_PROCESS_PORT:
'8476',
xenv.TPU_PROCESS_ADDRESSES:
'localhost:8476,localhost:8477,localhost:8478,localhost:8479',
xenv.TPU_VISIBLE_CHIPS:
'0',
}),
('v4-8_process_0', {
'ACCELERATOR_TYPE': 'v4-8',
xenv.TPU_PROCESS_BOUNDS: '1,1,1',
Expand Down
2 changes: 1 addition & 1 deletion torch_xla/experimental/tpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,7 +172,7 @@ def configure_topology(local_rank: int,
tpu_env = get_tpu_env()

accelerator_type = tpu_env[xenv.ACCELERATOR_TYPE]
if version() == 4:
if version() >= 4:
# Process bounds with 4 chips per process
default_process_bounds = MeshShape.from_string(
tpu_env[xenv.TPU_PROCESS_BOUNDS])
Expand Down