Fix negative offsets in GPUBiasingMultiModel.remove_model

nemo/collections/asr/parts/context_biasing/biasing_multi_model.py

@@ -562,22 +562,19 @@ def remove_model(self, model_id: int):
         self.model2num_states[model_id] = 0
         self.model2num_arcs[model_id] = 0
         self.model2num_arcs_extended[model_id] = 0
-        # shift model offsets
+        # Shift the offsets of the models that lived after the removed one.
+        # Only ACTIVE models take part in the shift. An unused slot holds offset 0, and 0 is also a
+        # legitimate arena address, so a shift applied to every slot underflows every unused slot to a
+        # negative offset whenever the removed model started at 0. `model2active[model_id]` is already
+        # False here, so the removed model is excluded from the mask and its own offsets are cleared last.
+        states_shift_mask = self.model2active & (self.model2states_offset >= start_state)
+        arcs_shift_mask = self.model2active & (self.model2arcs_offset >= start_arc)
+        self.model2states_offset[states_shift_mask] -= num_states
+        self.model2arcs_offset[arcs_shift_mask] -= num_arcs
+
+        # clear the removed model's own offsets last
         self.model2states_offset[model_id] = 0
         self.model2arcs_offset[model_id] = 0
-        # shift states and arcs offsets
-        torch.where(
-            self.model2states_offset < start_state,
-            self.model2states_offset,
-            self.model2states_offset - num_states,
-            out=self.model2states_offset,
-        )
-        torch.where(
-            self.model2arcs_offset < start_arc,
-            self.model2arcs_offset,
-            self.model2arcs_offset - num_arcs,
-            out=self.model2arcs_offset,
-        )
 
     def get_init_states(self, batch_size: int, bos=True) -> torch.Tensor:
         """

tests/collections/asr/decoding/test_biasing_multi_model.py

@@ -143,6 +143,67 @@ def test_add_then_remove_model(self, stt_en_conformer_transducer_small, device:
         assert multi_model.model2states_offset[model_id2].item() == 0
         assert multi_model.model2arcs_offset[model_id2].item() == 0
 
+        # Every offset must stay a valid, non-negative arena index, including the reserved slots that
+        # have never held a model. Before the ordering fix in remove_model, removing the model that
+        # started at offset 0 underflowed every zero-valued slot to -num_states / -num_arcs.
+        assert torch.all(multi_model.model2states_offset >= 0)
+        assert torch.all(multi_model.model2arcs_offset >= 0)
+
+        # Inactive slots carry a canonical, empty descriptor.
+        inactive = ~multi_model.model2active
+        assert torch.all(multi_model.model2states_offset[inactive] == 0)
+        assert torch.all(multi_model.model2arcs_offset[inactive] == 0)
+        assert torch.all(multi_model.model2num_states[inactive] == 0)
+        assert torch.all(multi_model.model2num_arcs_extended[inactive] == 0)
+
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("remove_index", [0, 1, 2])
+    def test_offsets_stay_non_negative_for_any_removal_position(
+        self, stt_en_conformer_transducer_small, device: torch.device, remove_index: int
+    ):
+        """Removing a model from any position leaves every offset a valid arena index.
+
+        Removing the model at the front of the arena is the interesting case: its start offset is 0,
+        which every unused reserved slot also holds, so a shift that does not exclude inactive slots
+        drives all of them negative.
+        """
+        tokenizer = stt_en_conformer_transducer_small.tokenizer
+        vocab_size = tokenizer.vocab_size
+
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size).to(device)
+        phrase_lists = [["alpha", "beta"], ["gamma"], ["delta", "epsilon", "zeta"]]
+        model_ids = [
+            multi_model.add_model(create_boosting_model(phrases, tokenizer, device), alpha=1.0)
+            for phrases in phrase_lists
+        ]
+
+        # A surviving model's scores must be unchanged by an unrelated removal.
+        survivor = model_ids[(remove_index + 1) % len(model_ids)]
+        states = multi_model.get_init_states(batch_size=2, bos=True)
+        survivor_ids = torch.full((2,), survivor, dtype=torch.long, device=device)
+        scores_before, _ = multi_model.advance(states=states, model_ids=survivor_ids)
+        scores_before = scores_before.clone()
+
+        multi_model.remove_model(model_ids[remove_index])
+
+        assert torch.all(multi_model.model2states_offset >= 0)
+        assert torch.all(multi_model.model2arcs_offset >= 0)
+
+        inactive = ~multi_model.model2active
+        assert torch.all(multi_model.model2states_offset[inactive] == 0)
+        assert torch.all(multi_model.model2arcs_offset[inactive] == 0)
+
+        active = multi_model.model2active
+        assert torch.all(
+            multi_model.model2states_offset[active] + multi_model.model2num_states[active]
+            <= multi_model.num_states_total
+        )
+
+        scores_after, _ = multi_model.advance(states=states, model_ids=survivor_ids)
+        assert torch.equal(scores_before, scores_after)
+
     @pytest.mark.unit
     @pytest.mark.with_downloads
     @pytest.mark.parametrize("device", DEVICES)