diff --git a/fbgemm_gpu/test/tbe/inference/nbit_forward_test.py b/fbgemm_gpu/test/tbe/inference/nbit_forward_test.py index 9c028af7b9..3d7997e2cb 100644 --- a/fbgemm_gpu/test/tbe/inference/nbit_forward_test.py +++ b/fbgemm_gpu/test/tbe/inference/nbit_forward_test.py @@ -687,6 +687,156 @@ def test_nbit_forward_cpu_multi_table_out_of_bounds_all_tables_defined( self.assertEqual(tuple(output.shape), (2, sum(Ds))) self.assertTrue(torch.isfinite(output.float()).all().item()) + def _execute_cpu_pooled_zero_fill( + self, + weights_ty: SparseType, + D: int, + output_dtype: SparseType, + Ls: list[int], + prune_positions: set[int], + weighted: bool, + ) -> None: + """Run a single-table CPU SUM-pooled nbit forward with every embedding row + set to 1.0, so each output dim equals the number of non-pruned indices in + the bag (unweighted) or the sum of their per-sample weights (weighted). + + Empty bags (L == 0) and fully-pruned bags must therefore be exactly 0. The + pooled CPU kernel zeroes those rows itself (each bag is accumulated into a + zeroed scratch buffer, then written to every output element), so this + asserts the behavior end-to-end without relying on any external pre-zero of + the output tensor. + """ + E = 100 + D_alignment = ( + 1 if weights_ty.bit_rate() % 8 == 0 else int(8 / weights_ty.bit_rate()) + ) + D = round_up(D, D_alignment) + + op = IntNBitTableBatchedEmbeddingBagsCodegen( + embedding_specs=[("", E, D, weights_ty, EmbeddingLocation.HOST)], + output_dtype=output_dtype, + pooling_mode=PoolingMode.SUM, + device="cpu", + ) + op.fill_random_weights() + + # Overwrite every row with 1.0 (dequant of 1.0 is exact for these types). + quant_weights, quant_scale_shift = quantize_embs(torch.ones(E, D), weights_ty) + weights, scale_shift = op.split_embedding_weights()[0] + weights.copy_(quant_weights) + if quant_scale_shift is not None: + self.assertIsNotNone(scale_shift) + scale_shift.copy_(quant_scale_shift) + + B = len(Ls) + total = sum(Ls) + offsets = torch.tensor([0, *np.cumsum(Ls)], dtype=torch.int) + indices = torch.randint(0, E, (total,), dtype=torch.int) + per_sample_weights = ( + torch.arange(total, dtype=torch.float) if weighted else None + ) + + prune_mask = torch.zeros(total, dtype=torch.bool) + for p in prune_positions: + prune_mask[p] = True + pruned_indices = indices.clone() + pruned_indices[prune_mask] = -1 + + kept = (~prune_mask).float() + if weighted: + assert per_sample_weights is not None + contrib = per_sample_weights * kept + else: + contrib = kept + expected = torch.zeros(B, D, dtype=torch.float) + for b in range(B): + start, end = int(offsets[b]), int(offsets[b + 1]) + expected[b, :] = contrib[start:end].sum() + + output = op( + indices=pruned_indices, + offsets=offsets, + per_sample_weights=per_sample_weights, + ) + + torch.testing.assert_close( + output.float().cpu(), + expected, + atol=1.0e-2, + rtol=1.0e-2, + equal_nan=False, # uninitialized data leaking as NaN must fail + ) + + # The -1 "pruned index" sentinel is only honored by the quantized + # (int8/int4) nbit SpMDM kernels, which explicitly skip idx == -1. The + # full-precision (fp32/fp16/bf16) kernel has no such case: it treats any + # idx < 0 as an out-of-range failure, aborting the whole table rather than + # skipping that one entry. Scoping pruning to the types that support it + # avoids exercising an unsupported combination. + _PRUNE_SENTINEL_SUPPORTED: frozenset[SparseType] = frozenset( + {SparseType.INT8, SparseType.INT4} + ) + + def test_nbit_forward_cpu_empty_and_pruned_bags_zero_fill(self) -> None: + """Pooled CPU nbit forward zeroes empty and fully-pruned bags itself. + + Ragged bag lengths include empty bags (L == 0) at the start, middle, and + end -- exercised for every weight type below, since an empty bag needs + no -1 sentinel and the kernel must zero it regardless of weight type. + Quantized (int8/int4) weight types additionally get ~1/3 of their + indices pruned to -1 (including one non-empty bag that is fully + pruned), since only their kernels honor that sentinel (see + _PRUNE_SENTINEL_SUPPORTED). Because every embedding row is 1.0, each + output dim must equal the count / weight-sum of non-pruned indices, and + empty/pruned bags must be exactly 0 -- verifying the kernel writes every + output row. + """ + Ls = [0, 1, 0, 3, 4, 0, 7, 8] + total = sum(Ls) + all_prune_positions = {i for i in range(total) if i % 3 == 0} + cases = [ + (SparseType.INT8, 24, SparseType.FP32), + (SparseType.INT4, 24, SparseType.FP16), + (SparseType.FP16, 32, SparseType.BF16), + ] + for weights_ty, D, output_dtype in cases: + prune_positions = ( + all_prune_positions + if weights_ty in self._PRUNE_SENTINEL_SUPPORTED + else set() + ) + for weighted in (False, True): + with self.subTest( + weights_ty=weights_ty, + output_dtype=output_dtype, + weighted=weighted, + ): + self._execute_cpu_pooled_zero_fill( + weights_ty, D, output_dtype, Ls, prune_positions, weighted + ) + + def test_nbit_forward_cpu_out_of_bounds_index_stays_defined(self) -> None: + """An out-of-bounds index is sanitized by the CPU bounds check (clamped to + row 0) before the kernel runs, so the op still returns a fully-written, + finite output -- there is no path that leaves the (no-longer-pre-zeroed) + pooled output partially uninitialized.""" + E, D = 100, 16 + op = IntNBitTableBatchedEmbeddingBagsCodegen( + embedding_specs=[("", E, D, SparseType.INT8, EmbeddingLocation.HOST)], + output_dtype=SparseType.FP16, + pooling_mode=PoolingMode.SUM, + device="cpu", + ) + op.fill_random_weights() + + # bag 0: in-range index; bag 1: out-of-range index (clamped to row 0). + offsets = torch.tensor([0, 1, 2], dtype=torch.int) + indices = torch.tensor([0, E + 5], dtype=torch.int) # E + 5 is out of range + output = op(indices=indices, offsets=offsets) + + self.assertEqual(tuple(output.shape), (2, D)) + self.assertTrue(torch.isfinite(output.float()).all().item()) + @unittest.skipIf(*gpu_unavailable) @given( T=st.integers(min_value=1, max_value=10),