From 0f886a24b21dcc0a55c9d223cead6f45c35ae9eb Mon Sep 17 00:00:00 2001 From: sambhavnoobcoder Date: Fri, 3 Oct 2025 16:02:43 +0530 Subject: [PATCH 1/6] init --- .../models/switch_transformers/modeling_switch_transformers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/switch_transformers/modeling_switch_transformers.py b/src/transformers/models/switch_transformers/modeling_switch_transformers.py index 935152b4ff49..be1af855d83c 100644 --- a/src/transformers/models/switch_transformers/modeling_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modeling_switch_transformers.py @@ -898,7 +898,7 @@ def _prepare_4d_causal_attention_mask_with_cache_position( **kwargs, ): """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + Creates causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. Args: From 17a64125a438037af5b936e71dea90b7f0586451 Mon Sep 17 00:00:00 2001 From: sambhavnoobcoder Date: Fri, 3 Oct 2025 16:15:09 +0530 Subject: [PATCH 2/6] jitter-noise changes copied here --- .../modeling_switch_transformers.py | 16 +++++-- .../test_modeling_switch_transformers.py | 47 +++++++++++++++++++ 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/src/transformers/models/switch_transformers/modeling_switch_transformers.py b/src/transformers/models/switch_transformers/modeling_switch_transformers.py index be1af855d83c..b5293917ba0d 100644 --- a/src/transformers/models/switch_transformers/modeling_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modeling_switch_transformers.py @@ -102,11 +102,17 @@ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tens # https://huggingface.co/papers/2101.03961. # We also store the previous dtype to cast back the output to the previous dtype self.input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(self.dtype) + + # Create a copy for applying jitter noise + routing_states = hidden_states.clone() + routing_states = routing_states.to(self.dtype) + if self.training and self.jitter_noise > 0: - # Multiply the token inputs by the uniform distribution - adding some noise - hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) - router_logits = self.classifier(hidden_states) + # Apply jitter noise only to the routing copy + routing_states *= torch.empty_like(routing_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) + + # Use jittered states for routing decisions + router_logits = self.classifier(routing_states) # Apply Softmax and cast back to the original `dtype` router_probs = nn.functional.softmax(router_logits, dim=-1, dtype=self.dtype).to(self.input_dtype) @@ -898,7 +904,7 @@ def _prepare_4d_causal_attention_mask_with_cache_position( **kwargs, ): """ - Creates causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. Args: diff --git a/tests/models/switch_transformers/test_modeling_switch_transformers.py b/tests/models/switch_transformers/test_modeling_switch_transformers.py index 86238c053a35..f779439bbbf3 100644 --- a/tests/models/switch_transformers/test_modeling_switch_transformers.py +++ b/tests/models/switch_transformers/test_modeling_switch_transformers.py @@ -1024,6 +1024,53 @@ def test_max_routing_capacity(self): assert torch.sum(expert_index) <= batch_size * self.config.num_experts * self.config.expert_capacity + def test_jitter_noise_preserves_hidden_states(self): + r""" + Test that jitter noise is applied only to routing decisions and does not modify the original hidden states. + This tests the fix for the jitter noise issue where noise was corrupting the input hidden states. + """ + # Create a config with jitter noise enabled + config = SwitchTransformersConfig( + num_experts=2, + hidden_size=4, + d_ff=8, + router_jitter_noise=0.1, # Enable jitter noise + expert_capacity=4, + ) + + # Create router + router = SwitchTransformersTop1Router(config) + router.eval() # Set to eval mode first to test training mode separately + + # Create input hidden states + hidden_states = torch.tensor([ + [[0.5, 0.2, 0.1, 0.3], + [0.4, 0.6, 0.2, 0.8]] + ], dtype=torch.float32) + + # Test in eval mode - no jitter noise should be applied + original_hidden_states = hidden_states.clone() + with torch.no_grad(): + router_probs, expert_index, router_logits = router(hidden_states) + + # Hidden states should remain unchanged in eval mode + self.assertTrue(torch.equal(hidden_states, original_hidden_states)) + + # Test in training mode - jitter noise should be applied only internally + router.train() + torch.manual_seed(42) # Set seed for reproducible results + + original_hidden_states = hidden_states.clone() + with torch.no_grad(): + router_probs_train, expert_index_train, router_logits_train = router(hidden_states) + + # Hidden states should still remain unchanged after router call + self.assertTrue(torch.equal(hidden_states, original_hidden_states)) + + # Results should be different between eval and train mode due to jitter noise + # (though this might occasionally fail due to randomness, it's very unlikely with seed) + self.assertFalse(torch.allclose(router_logits, router_logits_train, atol=1e-5)) + @slow @require_torch From 83374dc8016e166bf7f28415a0f95b59c3cbb2b8 Mon Sep 17 00:00:00 2001 From: sambhavnoobcoder Date: Fri, 3 Oct 2025 16:26:44 +0530 Subject: [PATCH 3/6] ruff fix --- .../models/switch_transformers/modeling_switch_transformers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/switch_transformers/modeling_switch_transformers.py b/src/transformers/models/switch_transformers/modeling_switch_transformers.py index b5293917ba0d..346356e8056b 100644 --- a/src/transformers/models/switch_transformers/modeling_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modeling_switch_transformers.py @@ -629,7 +629,7 @@ def _init_weights(self, module): module.weight.data.fill_(factor * 1.0) elif isinstance( module, - (SwitchTransformersModel, SwitchTransformersForConditionalGeneration, SwitchTransformersEncoderModel), + SwitchTransformersModel | SwitchTransformersForConditionalGeneration | SwitchTransformersEncoderModel, ): module.shared.weight.data.normal_(mean=0.0, std=factor * 1.0) if hasattr(module, "lm_head") and not self.config.tie_word_embeddings: From 330723257d8e003c263ab6a22a65171836ddd0df Mon Sep 17 00:00:00 2001 From: sambhavnoobcoder Date: Fri, 3 Oct 2025 16:31:26 +0530 Subject: [PATCH 4/6] yes , another ruff one --- .../switch_transformers/modeling_switch_transformers.py | 4 +++- .../switch_transformers/test_modeling_switch_transformers.py | 5 +---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/transformers/models/switch_transformers/modeling_switch_transformers.py b/src/transformers/models/switch_transformers/modeling_switch_transformers.py index 346356e8056b..689e15535eb2 100644 --- a/src/transformers/models/switch_transformers/modeling_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modeling_switch_transformers.py @@ -109,7 +109,9 @@ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tens if self.training and self.jitter_noise > 0: # Apply jitter noise only to the routing copy - routing_states *= torch.empty_like(routing_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) + routing_states *= torch.empty_like(routing_states).uniform_( + 1.0 - self.jitter_noise, 1.0 + self.jitter_noise + ) # Use jittered states for routing decisions router_logits = self.classifier(routing_states) diff --git a/tests/models/switch_transformers/test_modeling_switch_transformers.py b/tests/models/switch_transformers/test_modeling_switch_transformers.py index f779439bbbf3..2a3da6931911 100644 --- a/tests/models/switch_transformers/test_modeling_switch_transformers.py +++ b/tests/models/switch_transformers/test_modeling_switch_transformers.py @@ -1043,10 +1043,7 @@ def test_jitter_noise_preserves_hidden_states(self): router.eval() # Set to eval mode first to test training mode separately # Create input hidden states - hidden_states = torch.tensor([ - [[0.5, 0.2, 0.1, 0.3], - [0.4, 0.6, 0.2, 0.8]] - ], dtype=torch.float32) + hidden_states = torch.tensor([[[0.5, 0.2, 0.1, 0.3], [0.4, 0.6, 0.2, 0.8]]], dtype=torch.float32) # Test in eval mode - no jitter noise should be applied original_hidden_states = hidden_states.clone() From 603fda28c43f2de8f2dc065f424a3d3f389a746e Mon Sep 17 00:00:00 2001 From: sambhavnoobcoder Date: Fri, 3 Oct 2025 16:47:52 +0530 Subject: [PATCH 5/6] modular fix --- .../modular_switch_transformers.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/transformers/models/switch_transformers/modular_switch_transformers.py b/src/transformers/models/switch_transformers/modular_switch_transformers.py index cf4eaf0cedff..ebc1fc77de1e 100644 --- a/src/transformers/models/switch_transformers/modular_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modular_switch_transformers.py @@ -159,11 +159,19 @@ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tens # https://huggingface.co/papers/2101.03961. # We also store the previous dtype to cast back the output to the previous dtype self.input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(self.dtype) + + # Create a copy for applying jitter noise + routing_states = hidden_states.clone() + routing_states = routing_states.to(self.dtype) + if self.training and self.jitter_noise > 0: - # Multiply the token inputs by the uniform distribution - adding some noise - hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) - router_logits = self.classifier(hidden_states) + # Apply jitter noise only to the routing copy + routing_states *= torch.empty_like(routing_states).uniform_( + 1.0 - self.jitter_noise, 1.0 + self.jitter_noise + ) + + # Use jittered states for routing decisions + router_logits = self.classifier(routing_states) # Apply Softmax and cast back to the original `dtype` router_probs = nn.functional.softmax(router_logits, dim=-1, dtype=self.dtype).to(self.input_dtype) From 7aa71733831f11328b3aa5ac77a2ac23872547c6 Mon Sep 17 00:00:00 2001 From: sambhavnoobcoder Date: Fri, 3 Oct 2025 16:53:32 +0530 Subject: [PATCH 6/6] modular fix --- .../models/switch_transformers/modular_switch_transformers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/switch_transformers/modular_switch_transformers.py b/src/transformers/models/switch_transformers/modular_switch_transformers.py index ebc1fc77de1e..ec18790f0940 100644 --- a/src/transformers/models/switch_transformers/modular_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modular_switch_transformers.py @@ -360,7 +360,7 @@ def _init_weights(self, module): module.weight.data.fill_(factor * 1.0) elif isinstance( module, - (SwitchTransformersModel, SwitchTransformersForConditionalGeneration, SwitchTransformersEncoderModel), + SwitchTransformersModel | SwitchTransformersForConditionalGeneration | SwitchTransformersEncoderModel, ): module.shared.weight.data.normal_(mean=0.0, std=factor * 1.0) if hasattr(module, "lm_head") and not self.config.tie_word_embeddings: