From 5ab31f4523de1c4df0d7194a2015dd856afd6ca4 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 20 Apr 2026 19:03:42 +0000 Subject: [PATCH] fix: use no_mask in local ViT layer spec The local ViT layer spec incorrectly used AttnMaskType.causal while the TE spec uses AttnMaskType.no_mask. ViTs use bidirectional self-attention over image patches, so causal masking is incorrect. Commit 20abc8599 ("fix vit mask") fixed this in the TE path but missed the local path. Fixes #1791 Co-Authored-By: Claude Opus 4.7 (1M context) --- megatron/core/models/vision/vit_layer_specs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/models/vision/vit_layer_specs.py b/megatron/core/models/vision/vit_layer_specs.py index 51074b3836e..e5b08855846 100644 --- a/megatron/core/models/vision/vit_layer_specs.py +++ b/megatron/core/models/vision/vit_layer_specs.py @@ -69,7 +69,7 @@ def get_vit_layer_with_local_spec() -> ModuleSpec: input_layernorm=LNImpl, self_attention=ModuleSpec( module=SelfAttention, - params={"attn_mask_type": AttnMaskType.causal}, + params={"attn_mask_type": AttnMaskType.no_mask}, submodules=SelfAttentionSubmodules( linear_qkv=ColumnParallelLinear, core_attention=DotProductAttention,