Skip to content

Improve Mosaic memory profiling tutorial #3753

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

Jump to bottom
Open
basilwong wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Improve Mosaic memory profiling tutorial #3753

basilwong wants to merge 4 commits into from

Conversation

@basilwong
Copy link
Contributor

@basilwong basilwong commented Jan 28, 2026 •
edited
Loading

Summary

This PR improves the Mosaic memory profiling tutorial with several enhancements:

  • Add GPT-2 memory profile images: Visual comparison of memory usage with and without activation checkpointing
  • Add Google Colab download instructions: Code block showing how to download generated snapshot/profile files
  • Fix subprocess output visibility: Mosaic CLI output is now captured and printed so users can see the analysis results
  • Split Mosaic analysis into separate code blocks: Better readability for the baseline vs buggy model comparison
  • Refactor GPT2WithDebugOverhead: Changed from subclassing to wrapper pattern, fixing transformers version compatibility
  • Update section formatting: Bold headers instead of RST section underlines for cleaner rendering

Test plan

  • Build tutorial locally with GALLERY_PATTERN="mosaic_memory_profiling_tutorial" make html-noplot
  • Verify new images render correctly
  • Verify subprocess output is visible in rendered tutorial
  • Verify GPT2WithDebugOverhead works with current transformers version

@basilwong
Improve Mosaic tutorial: add images, fix subprocess output, refactor …
d745ba2 
…buggy model

- Add GPT-2 memory profiling images (with/without activation checkpointing)
- Add Google Colab download instructions for generated files
- Fix subprocess.run to capture and print Mosaic CLI output
- Split Mosaic analysis into separate code blocks for readability
- Refactor GPT2WithDebugOverhead to use wrapper pattern instead of
  subclassing, fixing transformers version compatibility issues
- Remove try/except workaround that was bypassing the tutorial's purpose
- Update section formatting (bold headers instead of RST underlines)
@pytorch-bot
Copy link

pytorch-bot bot commented Jan 28, 2026 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/tutorials/3753

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit 5571bd9 with merge base b99753a (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla bot added the cla signed label Jan 28, 2026
@basilwong basilwong mentioned this pull request Jan 28, 2026
Closed
svekars
svekars reviewed Jan 28, 2026
View reviewed changes
.ci/docker/requirements.txt
evaluate
accelerate>=0.20.1
git+https://github.com/facebookresearch/mosaic.git
altair
Copy link
Contributor

@svekars svekars Jan 28, 2026

Choose a reason for hiding this comment

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

do you need omegaconf too? https://docs-preview.pytorch.org/pytorch/tutorials/3753/beginner/mosaic_memory_profiling_tutorial.html#analyze-the-baseline-clean-snapshot

@basilwong basilwong force-pushed the mosaic-tutorial-improvements branch from 25d498d to f26fe03 Compare January 28, 2026 21:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@svekars svekars svekars approved these changes

Assignees

No one assigned

Labels

cla signed

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@basilwong @svekars