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.

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

Update code and README #122

Merged
merged 2 commits into from
Jun 3, 2024
Merged

Conversation

yanbing-j
Copy link
Collaborator

This PR is to update README of PyTorch nightly version, and update some codes related to CPU support.

We have optimized attn_mask in SDPA with block-wise algorithm in pytorch/pytorch#126961. It can reduce memory peak usage, and have some performance improvements as below.

Performance

Single socket in Intel (R) Xeon (R) CPU Max 9480
Dtype: bfloat16, models: vit_b and vit_h, test in SDPA and Triton commit https://github.com/pytorch-labs/segment-anything-fast/blob/main/experiments/run_experiments.py#L199-L200, select the time of 20th iteration.

  vit_b   vit_h  
  attn_mask w/o block-wise attn_mask w/ block-wise attn_mask w/o block-wise attn_mask w/ block-wise
SDPA 10.95s/it 6.59s/it 19.93s/it 12.33s/it
Triton 10.66s/it 7.12s/it 19.87s/it 12.26s/it

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 31, 2024
@yanbing-j
Copy link
Collaborator Author

Hi @cpuhrsch , could you please take a look? Thanks!

total_memory = psutil.virtual_memory().total
max_memory_allocated_bytes = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
max_memory_allocated_percentage = int(100 * (max_memory_allocated_bytes / (total_memory >> 10)))
max_memory_allocated_bytes = max_memory_allocated_bytes >> 10

if print_header:
print(",".join(["sam_model_type", "batch_size", "memory(MiB)", "memory(%)", "img_s(avg)", "batch_ms(avg)/batch_size", "mIoU", "use_compile",
Copy link
Contributor

Choose a reason for hiding this comment

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

Is there something you could add to indicate the device? Seems like this needs a new column to track the device used.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Yes, of course. I just add device at the front.

Copy link
Contributor

@cpuhrsch cpuhrsch left a comment

Choose a reason for hiding this comment

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

Please see comment

@cpuhrsch cpuhrsch merged commit 3e9c47d into pytorch-labs:main Jun 3, 2024
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants