Repository navigation
Conversation
Carry the MiniMax-M3.1 GB300 serving work onto this branch as one change. Each item is gated by an environment variable or config field and is off by default, except where noted. - NVFP4 index-K (indexer_kv_dtype='nvfp4'): the NVFP4 horizontal producer emits packed NVFP4 index-K, the V2 cache manager serves the (data, scale) views, decode and verify score them with the MiniMax MSA-v1 Q8KV4 indexer, and prefill dequantizes the context rows' pages to E4M3. - MXFP8 boundary quant: the boundary AllReduce emits the next layer's MXFP8 QKV input (always on where it applies; MNNVL and unaligned hidden sizes keep the plain fusion). - KV cache manager v2: host-tier copies staged one step ahead, and layer-wise host-KV onboarding. - AllReduce fusions: an AllReduce+RMSNorm epilogue emitting NVFP4 in linear SF layout, and the MoE-boundary AllReduce fusions (TRTLLM_ARIA_AR_NVFP4, TRTLLM_ARIA_MOE_AR). - The never-mix scheduling, sampler, fusion and prefill stack. - MXFP8 decode GEMM: an opt-in fixed swap-AB CuTeDSL tactic for 33-256 rows (TRTLLM_ARIA_LOWM_GEMM). - aria/m31f: recipes, env switches, build and run notes. Where this work and the Rubin (sm_107) commits change the same code, both are kept. NVFP4 index-K has two writer/reader pairs over the same 72-byte-per-token pages: the built-in model's Q8KV4 path above, and the sm_107 targets' FP8-sink writer with their own readers. arm_sm107_msa_paths now also sets cache_manager.NVFP4_INDEX_K_SINK, which selects the sm_107 pair, so every other process takes the Q8KV4 path. The sm_107 host-tier prefix prefetch and the staging above live side by side in PyExecutor and KVCacheManagerV2. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
… under disaggregated serving (NVIDIA#25) * [None][feat] Unified kv cache manager foundation (NVIDIA#20178) Signed-off-by: Allison Lim <allim@nvidia.com> (cherry picked from commit fce28f0) Backport to this branch: the cpp/ part (BufferConfig.coalescing_group, the attentionOp 3-D pool-pointer table) is not included; _coalescing_group_kwargs passes the group only when it is nonzero, so layer specs without NVFP4 buffers run on the existing bindings, and the two mixed NVFP4/BF16 pool tests skip on bindings without coalescing_group. Hunks for main-only features (mixed-precision KV, pool rebalancing) are dropped, and layer_specs is declared on the class for managers built without __init__. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [TRTLLM-16722][feat] Dflash Unified KVCache (NVIDIA#20486) Port of NVIDIA#20486 at 219e26e (net diff against the main it merges, df95eda). Signed-off-by: Allison Lim <allim@nvidia.com> Backport to this branch: conflicts resolved against the separate draft pool's disagg window transfer (RankInfo keeps both draft fields), the drafter slot parking and seeds in DFlashSpecMetadata.prepare (the position seeding applies to the unified cache only), and the fused block-table refresh (skipped when the layer has a block offset). Hunks for code this branch does not have (spec_tail_kernels, model-level layer specs) are dropped; Gemma4 hybrid geometry stands in for _has_variable_kv_cache_geometry. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][feat] Per View KV head transfer (NVIDIA#20709) Port of NVIDIA#20709 at 300a671 (net diff against the main it merges, 0746a23). Signed-off-by: Allison Lim <allim@nvidia.com> Backport to this branch: import and registration-check conflicts resolved by keeping both sides. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Bound DFlash's shared context by the allocation's extent DFlash took the number of valid block-table entries as the pages a request holds. On a unified cache with per-layer page tables (e.g. an FP8 target with a BF16 drafter), a sliding-window draft layer's evicted pages stay BAD_PAGE_INDEX in front of its live tail, so the count stopped at about one window (66 blocks for window 2048, page 32) whatever the context length. The decode-time allocation check then trapped once a context passed ~2.1K tokens, and prefill skipped the drafter's context store for later chunks and for prefix hits past the window. Use the extent of the allocation instead, the last valid entry + 1. Reads below the window are never issued and writes land in the live tail; for a full-attention drafter, and on the cpp backend, the extent equals the count. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> (cherry picked from commit ff11364) * [None][fix] Re-seed DFlash generation slots in CUDA-graph warmup CUDA-graph warmup builds every batch from dummy requests whose ids repeat across batches (0..bs-1), so a reused id keeps its drafter slot and the previous batch's context length, which the replay after each capture advances for good. A later batch can give the same id far fewer pages, and on a unified KV cache the drafter's allocation check then trapped at server start: during the advanced-sampling capture pass the filler rows reached 31..79 positions against a 24-position allocation. Only warmup runs a CUDA-graph batch eagerly, so there re-seed each generation row's slot from this batch's positions, as prepare() already seeds a new request. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> (cherry picked from commit 2b23449) * [None][feat] MiniMax-M3.1 DSpark: keep the drafter's KV in the target's manager under disagg On GB300 a disaggregated generation worker gets the target's prompt KV but not the DSpark drafter's, which lives in a separate pool: the drafter starts from an empty window at block 0 and acceptance collapses on long prompts. Hosting the draft layers in the target's KVCacheManagerV2 lets the transceiver move them with the target's pages. - MiniMaxM3KVCacheManagerV2 hosts explicit draft layer specs: index-K buffers follow the spec-built ones, dense target layers stay FP8, size-only coalescing keeps the target pools, draft layers use the base view. - KvCacheCreator unifies a standalone DSpark drafter with such a manager under disaggregated serving when the drafter reads the target's page size. Agg keeps the separate pool paired with block reuse, and a model that ships that pool's window (F0a) keeps it in disagg. TRTLLM_DISAGG_SHARED_DRAFT_KV=0 opts out. - An NVFP4 target transfers with a BF16 draft held in pools of its own. - A windowed draft lifecycle gets the hot-tier bytes of every admissible sequence's window plus one context chunk, instead of the bytes-per-token split that handed it a quarter of the KV cache. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][feat] Reuse a unified DSpark drafter's KV only where its window is complete With the drafter's layers in the target's KV cache manager, the draft lifecycle's REQUIRED policy let the drafter's window constrain every prefix match: a request branching off a shared prefix found no committed window there and re-prefilled its whole prompt. - A windowed drafter's layer specs are OPTIONAL. The target alone sets the match, and a new cache's first resume selects the drafter's layers where their window is complete (page-aligned matches only). - Where the window is missing, prepare_context caps the claim so the prefill writes everything the drafter reads (from L + 2 - W on): at L + 1 - W, or one window before a branch point within W of the prompt end, whose window the request then publishes. A partial match first retries at its last full page. - The commit gate for unstored drafter context, its unwritten-prefix marks and branch-point publication cover a manager that hosts the drafter's layers, not only a paired draft pool. A disagg generation request needs no cap: the KV transfer requests an unselected OPTIONAL group in full. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][perf] MiniMax-M3.1 DSpark: copy only the unified drafter's block offsets per step A hybrid NVFP4 MiniMax-M3 cache is pool-less: no target layer reads its attention-op block offsets. With explicit draft layers it kept every layer's pool, so each step converted the whole per-layer table (layers x sequence slots x max_seq_len / tokens_per_block) on the device: about 3 ms per decode step at max_seq_len 1M on TP2. Keep one attention-op pool, the first draft layer's, which is the only one the unified drafter reads, and fill it from the draft lifecycle's rows with the copy a separate draft pool uses. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Allocate the whole block of a drafter that shares the target's KV cache A unified DFlash/DSpark drafter appends its block after the request's context, through the request's own pages. The manager sized every request with num_extra_kv_tokens = K-1 beyond its history, but the DSpark shift_label block is K wide (DFlash K+1). When the allocation ended on a page boundary, at the prompt end or after a step that accepts all K+1 tokens, the block's last slot had no page: the drafter then truncated the advertised context and its scratch overwrote the last context position in the draft layers, silently. When DRAFT-owner layers live in the manager, size each request beyond its history by max(num_extra_kv_tokens, drafter block width) at every allocation site and in the generation headroom. The creator passes the checkpoint-resolved width (shift_label). Adapted from NVIDIA/TensorRT-LLM Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> NVIDIA#20664 (76c4cef); num_extra_kv_tokens keeps its position-side meaning. * [None][fix] Fail at startup when TRTLLM_DISAGG_SHARED_DRAFT_KV=1 cannot be honored An explicit TRTLLM_DISAGG_SHARED_DRAFT_KV=1 on a disaggregated worker whose target and drafter cannot share one KV cache manager silently kept the separate draft pool, whose KV is not transferred. Raise at KV cache creation instead, as NVIDIA/TensorRT-LLM NVIDIA#20664's _check_shared_draft_kv_request does. disaggregated_mpi_worker now writes "auto" rather than "1" when both sections pass its coarse checks, so its automatic choice can still fall back where a worker's finer checks refuse; the creator shares under "1" and "auto" alike. The unset default is unchanged. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Let disaggregated_mpi_worker share a DSpark drafter's KV The launcher wrote TRTLLM_DISAGG_SHARED_DRAFT_KV=0 for any drafter other than DFlash, and for any page size other than 16, 32 or 64, so a MiniMax-M3.1 DSpark deployment started through it (128-token pages) kept the separate draft pool, whose KV is not transferred. Accept a standalone DSpark drafter with attention_backend unset, AUTO or TRTLLM, as NVIDIA/TensorRT-LLM NVIDIA#20664 does, and leave its page size to each worker's KvCacheCreator: the launcher cannot know whether the drafter reads the target's page. It writes "auto"; if the two sides still disagree, the peers' draft-layer check fails loudly. An explicit TRTLLM_DISAGG_SHARED_DRAFT_KV=0 is now kept. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Keep attention-DP dummy generation rows on DFlash's dummy slot A generation row with the attention-DP dummy id must never inherit a real slot, because the C++ offset copy maps placeholders to page 0: a slot that a context warmup stored under that id would let the dummy row write its drafter K/V past its own pages, onto another request's page. As in NVIDIA/TensorRT-LLM NVIDIA#20664, such rows now always take the dummy slot. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][perf] DFlash: size a unified MiniMax-M3.1 drafter's context stores on the host DFlash bounds each context store by the row's allocation, which it reads back from the device unless the drafter's block-table rows are known to be written in full. That was only recognized for a separate draft pool, so with the unified drafter every context chunk waited for the target's forward before the store was issued. Let the KV cache manager say whether it rewrites copied rows in full, and check the target manager's own table for a unified drafter. The MiniMax-M3 manager's drafter-only pool is copied in full. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][feat] KVCacheManagerV2: reuse an OPTIONAL group at a partial endpoint An OPTIONAL lifecycle group was checked only at a page-aligned endpoint, so a partial prefix match always left it incomplete. A unified DSpark drafter then lost its window's partial page and the caller cut its claim back to the last full page, recomputing up to a page per request. Check an OPTIONAL group at any endpoint with the REQUIRED coverage rule: the last page needs only the matched prefix. The first resume copies a selected group's partial page into a private page, as for REQUIRED, and counts it as a partial reuse. Partial snapshots now publish OPTIONAL pages as well, so a continuation of a prompt that ended mid-page finds them. A match whose partial page is missing still falls back to the last full page. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Give the block-table refresh test's stub the extent columns The stock DFlashWorker._refresh_ctx_block_tables reads _ctx_block_cols since the allocation-extent fix; the test's object.__new__ stub never set it. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> --------- Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> Co-authored-by: allisonlim-nv <allim@nvidia.com>
…'s manager in agg too (NVIDIA#26) Aggregated serving kept the DSpark drafter's KV in a separate pool paired with the target's block reuse, while disaggregated serving keeps it in the target's KVCacheManagerV2 as a windowed draft lifecycle. Use that layout in both, so one mechanism serves agg and disagg. What the paired pool did carries over: the drafter's window is reused only where it is complete, unwritten drafter pages are never published, and the draft lifecycle's share is window-sized for the same admissible sequences (the overlap scheduler's two per slot in agg). A model that ships the separate pool's window itself (F0a) keeps that pool in both modes. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…s for the request lifetime (NVIDIA#27) * [None][perf] DSpark: temperature-only draft sampling runs no filter DSpark's ancestral drafting samples every block position with the temperature only, but passed neutral top-k / top-p / min-p tensors, so under advanced_sampling_mode no_topk each of the K positions still ran a top-p renorm (a keep-all pass at top_p = 1). Pass None instead: the FlashInfer chain then runs softmax and sampling only, and the fused op (mode full) receives its disable sentinels as before. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][perf] DSpark: draw ancestral drafts with one fused kernel per position Ancestral DSpark drafting ran a full sampling pipeline for every block position (bias add, softmax, sampling, RNG bookkeeping and a [B, vocab] proposal scatter), where the greedy chain it replaced ran one batched sampling over the B x K rows. Draw position k by Gumbel-max in one Triton kernel instead: it adds the Markov bias, scales by 1/T and takes the argmax of the scaled logits plus Gumbel noise from a Philox stream keyed by the request's (seed, offset) draft slot of position k, so every TP rank and every graph replay draws the same token. It writes the biased logits in float32, whatever the drafter's logits dtype, and per-chunk softmax statistics; a small kernel picks the token and gathers the next Markov embedding (the bias GEMM stays on cuBLAS). After the block, one kernel writes q_k = softmax(logits_k / T) for all positions into the slot-indexed draft_probs from those float32 logits, so the stored proposal is the distribution each token was drawn from. CUDA-graph padding rows draw from the bias alone. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][perf] Keep a DFlash request's drafter slot while it sits out an iteration DFlashSpecMetadata.prepare freed the drafter slot of every request missing from the batch, so each context-only step (ARIA NO_MIX) evicted every generation request and the next decode step paid a slot clear and a parked context resume per request. A slot now stays with its request: prepare keeps the slot map in batch order, and _assign_slot reclaims the least recently batched holder only when no slot is free, parking it as the eviction did. A slot that leaves its request (reclaim or restarted prefill) is reset at once: with free slots now rare, a caller's deferred clear would hand it back out with its old base pages. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…kers that ship its window (NVIDIA#24) * [None][fix] Keep the standalone drafter's pool unpaired on disagg workers that ship its window NVIDIA#22 pairs the standalone DSpark drafter's pool with the target's block reuse whenever the two page sizes match. A disaggregated worker of a model that opts into the draft-window transfer (RuntimeHooks.disagg_draft_kv: the sm_107 MiniMax-M3.1 targets) ships an unpaired pool only, so with block reuse on, select_disagg_draft_kv_cache_manager refused the paired pool at executor creation and the worker failed to start. KvCacheCreator._joint_reuse_supported now declines the pairing for a disaggregated worker that ships the window, using the selector's own opt-in predicate (model_ships_disagg_draft_window). Aggregated serving keeps the pairing on every GPU; the kill switch TRTLLM_M31_DISABLE_DISAGG_DRAFT_KV=1 keeps it on disagg workers too. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Keep the static draft page on disagg workers that ship the drafter window The target's 128-token draft pages (NVIDIA#22) serve only the pairing with the target's block reuse. A disaggregated worker that ships the drafter window keeps its pool unpaired, and there a prefix hit's first stored position is rounded up to a draft page (DFlashWorker._ctx_base_for_store): a 128-token page hides up to 127 tokens of the tail from the drafter, a 64-token page up to 63. On Rubin 1P1D disagg (K7 block, 96% prefix hits) the 128-token unpaired pool lost ~0.16 AL against the 64-token one, so such a worker now keeps the static page size, as before NVIDIA#22. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> --------- Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
…time hooks test_estimation_temporarily_uses_inferred_pool_sizing builds KvCacheCreator on a bare Mock model engine. runtime_hooks_of now rejects a model whose runtime_hooks attribute is not a RuntimeHooks, and a Mock attribute is a Mock, so all five cases failed with "Mock.runtime_hooks must be a RuntimeHooks, got Mock" before reaching the code under test (on 519d279 as well). Set the attribute to None, as on a real model that declares no hooks (the all-off default). Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
…ult on Blackwell (KVLO_OFF) Layer-wise host-KV onboarding (TRTLLM_KV_LAYERWISE_ONBOARD, on by default since the ARIA serving commit) costs the VR TP2 aggregated K7 decode step 0.41 ms at 22 requests (D176 step synth on one node: 17.40 -> 16.99 ms with it off; an A/A/A control on one node drifts +0.3 / +0.6%), and the VR c56 agentic E2E is flat with it off (+0.8%). The choice moves KV pages only, so tokens are unchanged. MiniMax-M3.1 with DSpark on Blackwell-family GPUs (sm_100 / sm_103 / sm_107) now defaults to whole-migration waits. The new pyexecutor/kv_cache/m31_kv_defaults.arm_minimax_m31 arms it (kv_layerwise_onboard.arm_m31_layerwise_onboard_off) before the KV cache managers are built, from the sm_107 targets' routing, the sm_103 targets' routing and the built-in MiniMaxM3ForCausalLM with DSpark. Arming sets no environment variable: the shipped CuTe DSL AOT objects are keyed on the TRTLLM_M31_* switches. TRTLLM_KV_LAYERWISE_ONBOARD, when set, still wins (1 turns layer-wise onboarding on, 0 off); TRTLLM_M31_DISABLE_KVLO_OFF=1 keeps the layer-wise default. ruff-format also rewrapped four existing statements in modeling_minimaxm3.py (whitespace only). Note for GB300: this reverses the layer-wise onboarding default that the ARIA serving commit (519d279) introduced for GB300. It is measured on VR only; please validate it on GB300, and set TRTLLM_M31_DISABLE_KVLO_OFF=1 (or TRTLLM_KV_LAYERWISE_ONBOARD=1) to keep layer-wise onboarding there. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
…ault (UNPAIRED_DRAFT_POOL) Joint draft / target prefix reuse pairs the standalone DSpark drafter's KV pool with the target's prefix cache. On VR TP2 aggregated K7 serving the pairing costs decode time: - Decode step synth at 22 requests (D176, one node): 17.32 ms with the pairing, 16.78 ms without. - Agentic E2E, 3600 s, synthetic (forced) K7 AL 4.04 (natural AL not demonstrated by these runs), on this branch with vs without this change: X_P90 c68 161.0 vs 156.7, c56 195.2 vs 192.9, c32 (running-set cap 8) 322.6 vs 319.7, c20 487.4 vs 475.8 (+1 .. +3%); Y per GPU about +0.3%; prefix cache hit unchanged. - Natural AL, unforced: unchanged on a single-turn workload, where the reuse does not engage. On the agentic SA workload (c56, unforced), where it does, natural AL is about 1.8% lower with this change (pooled over three pairs: -0.076, 95% CI -0.126 .. -0.026), while natural-acceptance X_P90 is still +2.2% (Y flat). MiniMax-M3.1 with DSpark on Blackwell-family GPUs (sm_100 / sm_103 / sm_107) now keeps the separate draft pool unpaired by default: m31_kv_defaults.arm_minimax_m31 (kv_cache_manager_v2.arm_m31_unpaired_draft_pool) arms it before the KV cache managers are built, and KvCacheCreator._joint_reuse_supported declines the pairing. This applies to aggregated and disaggregated roles alike, wherever a separate draft pool remains: the sm_107 targets ship its window (RuntimeHooks disagg_draft_kv), so the unified target + drafter KV manager keeps their separate pool and pairs it in aggregated serving; where the unified manager takes the drafter, it is a no-op. Kill switch TRTLLM_M31_DISABLE_UNPAIRED_DRAFT_POOL=1 keeps the joint reuse pairing. Measured on VR only; please validate on GB300. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
… notes MSA-v1 sparse decode and NVFP4 index-K need MiniMax MSA built with the quad4 V-scale layout patch, which upstream MSA does not carry. Add the patch under m31_flash/msa/ with a README on applying it to MSA 1e7bf76f, and point aria/m31f/README.md at it. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…#20075) Apply upstream PR NVIDIA#20075 to the MiniMax-M3.1-Flash GB300 branch. Preserve the branch's runtime opt-in helper alongside the new request-state helper, and place named-allreduce topology warmup in DecoderRunner.warmup because this branch predates the standalone DecoderWarmup refactor. Retain upstream tests, apply the branch's formatter, and refresh copyright headers on modified sources. (cherry picked from commit 348347c) Signed-off-by: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com>
pcicotti
marked this pull request as draft
October 29, 2026 15:56
brb-nv
force-pushed
the
user/brb/m3.1-flash-gb300-initial
branch
from
October 30, 2026 20:23
7a44301 to
23c67f2
Compare
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
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.
Description
Backport NVIDIA/TensorRT-LLM#20075 (
348347ce3e932e4ac487ce7829d04f16f5b8d5bf) touser/brb/m3.1-flash-gb300-initial. This brings MXFP8 partial weight refit, torch.compile/piecewise CUDA graph refit handling, active-request KV recomputation, and Ray collective fixes to the MiniMax-M3.1-Flash GB300 branch.Resolve the two backport conflicts by retaining both the branch's
_arm_runtime_opt_inshelper and upstream's_request_flaghelper, and placingwarmup_named_allreduce_topologies(self.model)in the existingDecoderRunner.warmup(). This branch predates the separateDecoderWarmupcomponent introduced by NVIDIA#20134; its existing MSA warmup scope is preserved. Retain upstream test changes and apply the branch's formatting and copyright requirements.Test Coverage
pytest tests/unittest/_torch/executor/test_recompute_active_requests_disagg.py tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py -q; collection was blocked byModuleNotFoundError: No module named 'torch'on the local Mac.PR Checklist