Skip to content

Add limit to max inflight TPU computations#6533

Merged
will-cromar merged 2 commits intomasterfrom
wcromar/max-inflight-computations
Feb 14, 2024
Merged

Add limit to max inflight TPU computations#6533
will-cromar merged 2 commits intomasterfrom
wcromar/max-inflight-computations

Conversation

@will-cromar
Copy link
Copy Markdown
Collaborator

Requires XLA_REGISTER_INSTALLED_PLUGINS=1 to use DevicePlugin API.

def physical_chip_count(self):
return num_available_chips()

def client_create_options(self):
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

who calls this method?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The parent class' virtual method is pybinded and called during client creation here:

auto create_options = plugin->client_create_options();

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, Will. Hard to tell it it's a virtual method in pythong. haha

Copy link
Copy Markdown
Collaborator

@alanwaketan alanwaketan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM.

@will-cromar will-cromar merged commit 29b01c6 into master Feb 14, 2024
def client_create_options(self):
return {
'max_inflight_computations':
xu.getenv_as('XLA_TPU_MAX_INFLIGHT_COMPUTATIONS', int, 4)
Copy link
Copy Markdown
Collaborator

@miladm miladm Feb 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shall we track this new variable in https://github.com/pytorch/xla/blob/master/configuration.yaml @will-cromar?

cc @wonjoolee95

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants