fix: Disable torch.compile for PI05 policy on MPS device #2723
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Fixes InductorError when running PI05 policy on Apple Silicon (MPS). torch.compile on MPS does not support the welford_combine reduction operation used in layer normalization. Changes: - Check if device is MPS before applying torch.compile - Log a warning when compilation is skipped on MPS - Allow model to run without compilation (slower but functional) Error fixed: torch._inductor.exc.InductorError: NotImplementedError: welford_combine This enables PI05 policy inference on Apple Silicon Macs until full MPS support is added to torch.compile. Tested with: - MacBook Pro with Apple Silicon (MPS device) - PI05 policy with compile_model=True - SO100 Follower robot with dual camera setup
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
Fixes
InductorError: NotImplementedError: welford_combinewhen running PI05 policy inference on Apple Silicon (MPS device). This error occurs becausetorch.compileon MPS does not support the Welford reduction operation used in layer normalization.Problem
When running PI05 policy with
compile_model=Trueon Apple Silicon Macs:The error occurs during model compilation because the MPS backend doesn't implement the
welford_combinereduction used for computing mean and variance in layer normalization operations.Solution
Added a device check before applying
torch.compile:mps, skip compilation and log a warningTesting
Tested successfully with:
compile_model=Truein configRelated
This is a known limitation of PyTorch's MPS backend: pytorch/pytorch#150121