Add CPU fallback for triton EDT on macOS - #510
Open
msmackle wants to merge 14 commits into
Open
Conversation
triton is Linux/CUDA-only and fails to import on macOS. Wrap the import in try/except, guard the @triton.jit kernel behind _TRITON_AVAILABLE, and add a cv2.distanceTransform fallback in edt_triton for non-CUDA environments. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
|
Hi @msmackle! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
Falls back to CPU when CUDA is unavailable (e.g. macOS). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
meta-cla
Bot
added
the
CLA Signed
- decoder.py: coord cache precomputation during __init__ - vl_combiner.py: forward_text default device args - train/utils/distributed.py: all_gather device selection Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Add _DEVICE constant (cuda if available, else cpu) and replace all .cuda() tensor moves with .to(_DEVICE) throughout the file. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
aten::_addmm_activation is CUDA-only and raises NotImplementedError on MPS/CPU. Fall back to a standard matmul + activation when CUDA is unavailable. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
fused.py: use float32 in the CPU/MPS fallback path — MPS rejects bfloat16 accumulation in matrix multiplication and raises a failed assertion in MPSNDArrayMatrixMultiplication. io_utils.py: include MPS in _DEVICE selection so video frames are loaded onto the correct device on Apple Silicon. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Add _DTYPE constant (float16 on CUDA, float32 elsewhere) and replace all hardcoded float16/half() usages in loader functions. This ensures image tensors match the model's float32 dtype on MPS/CPU, preventing dtype mismatch errors in MPSNDArrayMatrixMultiplication. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Several inference-path functions unconditionally cast tensors to
bfloat16 for CUDA memory efficiency. On MPS/CPU this causes dtype
mismatches. Guard each cast with torch.cuda.is_available() so tensors
retain their existing dtype on non-CUDA devices.
Affected: sam3_tracking_predictor (maskmem_features x2),
sam3_image (backbone_fpn all-gather),
sam3_multiplex_detector (backbone_fpn all-gather x2).
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
The bfloat16 casts were applied unconditionally before the CUDA check, so mat1.dtype was always bfloat16 and the output was cast back to bfloat16 — causing dtype mismatches with float32 model weights. Restructure so bfloat16 casts only happen in the CUDA path. The non-CUDA path now casts inputs to the weight's dtype (float32 after model.float()), ensuring a consistent output dtype. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
sam3 has bfloat16 op incompatibilities with MPS. Since main.py forces DEVICE=cpu on non-CUDA systems, io_utils must also use cpu so image tensors and model weights land on the same device. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
On Apple Silicon, pin_memory() uses MPS-backed pinned memory, producing tensors on mps:0. Calling .to(device=cpu, non_blocking=True) on such a tensor then raises: RuntimeError: Attempted to set the storage of a tensor on device 'cpu' to a storage on different device 'mps:0' Wrap every pin_memory().to(device, non_blocking=True) pattern in a conditional so non-CUDA paths use a plain .to(device) instead. Files fixed: - sam3/model/geometry_encoders.py - sam3/model/sam3_multiplex_base.py - sam3/model/sam3_multiplex_tracking.py - sam3/model/sam3_tracker_base.py - sam3/model/sam3_video_inference.py - sam3/model/video_tracking_multiplex.py Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…n-CUDA On non-CUDA systems: - Skip torch.autocast(device_type='cuda') which is a no-op but entering its context manager can have unexpected side effects - Use storage_device='cpu' instead of hardcoded 'cuda' to prevent tensor device mismatches during tracker inference state initialization Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…bility Fixes RuntimeError on non-CUDA systems where .cuda() raises because CUDA is not available. Files: - sam3/model/sam3_tracker_base.py: maskmem features/pos_enc load - sam3/model/sam3_tracking_predictor.py: prev mask logits and image load Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
When the input tensor has batch_size=0 (no objects to process), the labels_list/counts_list stay empty and torch.stack([]) raises 'stack expects a non-empty TensorList'. This happens during text-mode propagation when a frame contains no detected instances. Early-return zero tensors of the correct shape when batch_size==0. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
triton is Linux/CUDA-only and fails to import on macOS. Wrap the import in try/except, guard the @triton.jit kernel behind _TRITON_AVAILABLE, and add a cv2.distanceTransform fallback in edt_triton for non-CUDA environments.