Repository navigation
fix attn_mask not updated error for informer model - #49038
kaixuanliu wants to merge 4 commits into
Conversation
Signed-off-by: kaixuanliu <kaixuan.liu@intel.com>
Signed-off-by: kaixuanliu <kaixuan.liu@intel.com>
| # Keep the 2D padding mask around: the distillation conv layers halve the sequence length, | ||
| # so the attention mask has to be downsampled together with the hidden states. | ||
| padding_mask = attention_mask if attention_mask is None or attention_mask.dim() == 2 else None | ||
| attention_mask = create_bidirectional_mask( |
There was a problem hiding this comment.
padding_mask = attention_mask if attention_mask is not None and attention_mask.dim() == 2 else None maybe so we keep only 2D masks
| def test_distillation_updates_encoder_attention_mask(self): | ||
| config, inputs_dict = self.model_tester.prepare_config_and_inputs() | ||
| config.distil = True | ||
| model = InformerModel(config).to(torch_device).eval() | ||
|
|
||
| transformer_inputs, _, _, _ = model.create_network_inputs(**inputs_dict) | ||
| encoder_inputs = transformer_inputs[:, : config.context_length] | ||
| attention_mask = torch.ones(encoder_inputs.shape[:2], device=torch_device, dtype=torch.long) | ||
|
|
||
| with torch.no_grad(): | ||
| outputs = model.encoder(inputs_embeds=encoder_inputs, attention_mask=attention_mask) | ||
|
|
||
| expected_sequence_length = (config.context_length + 1) // 2 | ||
| self.assertEqual( | ||
| outputs.last_hidden_state.shape[:2], | ||
| (self.model_tester.batch_size, expected_sequence_length), | ||
| ) | ||
|
|
There was a problem hiding this comment.
interesting, it would have failed loud with shape mismatch - do we have no tests with ConvLayers ?
|
[For maintainers] Suggested jobs to run (before merge) run-slow: informer |
CI recapDashboard: View test results in Grafana |
|
Sorry, forgot to reply. Seems like the slow CI ckpt has If it's just poor docs for old model, i'd say we don't need a fix, so could you verify it |
|
@zucchini-nlp , well I checked it through carefully again, it should be the same root cause as what is solved in #48975 , under compile mode the mask was wrongly created and lead to the error. After the fix the bug is gone, can close this PR now. |
Fix Informer encoder attention-mask handling when distillation is enabled.
The distillation convolution halves the encoder sequence length, but the attention mask previously retained its original length, causing a shape mismatch in later encoder layers under torch.compile.
The mask is now downsampled with the same pooling configuration as
InformerConvLayer.maxPooland rebuilt after every distillation step. pls help review, thx! @zucchini-nlp