From 18b184c929737966fad1f6c6c49ca14a7a996804 Mon Sep 17 00:00:00 2001 From: Huanyu He Date: Sat, 22 Jun 2024 08:33:41 -0700 Subject: [PATCH] implementation of fbgemm op - permute_multi_embedding (#2120) Summary: X-link: https://github.com/pytorch/FBGEMM/pull/2738 # context * current we have a working function `permute_pooled_embs_auto_grad` to do a full permute of KTs, including forward and backward * it has several limitations: a) it has to be a full permute, duplicates are not supported; b) in the main [use case](https://fburl.com/code/89od0rqm) there has to be a torch.concat on the input KTs, which is not very efficient; c) the function output a single KT which requires a split operation * there is some attempt to support duplicated outputs, but the backward doesn't work * this diff is trying to create a new kernel (named `permute_multi_embedding`) to support a multiple-KT to multiple-KT mapping operation with backward support # notes * this diff focuses on the implemenation and test of the operator * performance analysis and benchmark are in the next diff # operator example usage * used in python ``` # test inputs: 3 KTs with batch_size=2048 batch_size = 2048 keys = [["f1", "f2"], ["f3", "f4", "f5"], ["f6"]] lengths = [[96, 256], [512, 128, 768], [1024]] values = [ torch.randn(batch_size, sum(lens), device="cuda", requires_grad=True) for lens in lengths ] # target outputs: 4 KTs with re-arranged keys (features), duplicates are allowed groups = [["f1", "f3"], ["f2"], ["f4", "f1", "f6"], ["f1", "f5"]] # accessorial arguments to the op/kernel permutes, in_lengths, out_lengths = _multi_remap_to_groups( keys, lengths, groups ) # arguments outputs = torch.ops.fbgemm.permute_multi_embedding( values, permutes, in_lengths, out_lengths ) ``` * permutes ``` permutes = tensor( [ [0, 0, 0, 0, 3, 4], # f1 [1, 0, 0, 3, 5, 0], # f3 [0, 1, 3, 0, 4, 0], # f2 [1, 2, 5, 0, 6, 0], # f4 [0, 2, 0, 6, 3, -6], # f1 [2, 2, 0, 9, 8, 0], # f6 [0, 3, 0, 0, 3, -8], # f1 [1, 3, 11, 3, 7, 0], # f5 ] ) ``` # details 1. from the above example usage, we can clearly see that the operatior takes in the following: a) values: List[torch.Tensor], which represents the input KTs b) permutes: torch.Tensor, which contains the permute information, will be explained later c) output_lengths_list: List[int], the lengths of the output tensors (KTs), which is needed to allocate memory on device ahead d) in_lengths: torch.Tensor, lengths of input tensors, which is on device e) out_lengths: torch.Tensor, lengths of output tensors, which is on device 2. the operator returns a list of tensors, which represents the permuted KTs 3. `permute` is the most critical argument in this operator: a) 2-D tensor b) each row represents a key (feature) permute move c) a permute move = [input_tensor_id, output_tensor_id, input_start_idx, output_start_idx, feature_length, jump] d) jump is used in backward when a key (feature) from the input tensor is mapped to multiple places in the output tensors Differential Revision: D57055616 --- torchrec/sparse/jagged_tensor.py | 82 ++++++++++ torchrec/sparse/tests/test_jagged_tensor.py | 167 +++++++++++++++++++- 2 files changed, 247 insertions(+), 2 deletions(-) diff --git a/torchrec/sparse/jagged_tensor.py b/torchrec/sparse/jagged_tensor.py index f703fe8ec..a9e9e2e2c 100644 --- a/torchrec/sparse/jagged_tensor.py +++ b/torchrec/sparse/jagged_tensor.py @@ -36,6 +36,12 @@ torch.ops.load_library( "//deeplearning/fbgemm/fbgemm_gpu:permute_pooled_embedding_ops_cpu" ) + torch.ops.load_library( + "//deeplearning/fbgemm/fbgemm_gpu:permute_multi_embedding_ops_cpu" + ) + torch.ops.load_library( + "//deeplearning/fbgemm/fbgemm_gpu:permute_multi_embedding_ops_gpu" + ) except OSError: pass @@ -240,6 +246,82 @@ def _remap_to_groups( return permute, inv_permute, offsets, inv_offsets, splits +def _multi_remap_to_groups( + keys: List[List[str]], + key_lengths: List[List[int]], + groups: List[List[str]], +) -> Tuple[List[int], List[int], List[int]]: + """ + Given a list of keys and lengths per key for each group, return the permute 2D tensor, and 1D tensor lengths: + [[input_tensor_idx, output_tensor_idx, input_start, output_start, length]], [length] + """ + # key => (tensor_idx, key_index) + key_map: Dict[str, Tuple[int, int]] = { + key: (tensor_idx, key_idx) + for tensor_idx, tensor in enumerate(keys) + for key_idx, key in enumerate(tensor) + } + + # [offsets per tensor] + in_offsets: List[List[int]] = [[] for _ in key_lengths] + for i, tensor in enumerate(key_lengths): + in_offsets[i] = _cumsum(tensor) + in_lengths: List[int] = [sum(lengths) for lengths in key_lengths] + + # set total_permutes as the jump stop sign + total_permutes: int = sum(len(tensor) for tensor in groups) + out_lengths: List[int] = [0] * len(groups) + + # [input_tensor_idx, output_tensor_idx, input_start, output_start, length, jump] + permute_param = 6 + permutes: List[int] = [0] * (total_permutes * permute_param) + + # record the last seen index, so that can make the jump from last_seen to current + last_seen: Dict[str, int] = {} + permute_idx = 0 + for output_tensor_idx, output_tenser in enumerate(groups): + output_start = 0 + for output_key in output_tenser: + input_tensor_idx, input_key_idx = key_map[output_key] + input_start = in_offsets[input_tensor_idx][input_key_idx] + length = key_lengths[input_tensor_idx][input_key_idx] + + # add jump data + if output_key not in last_seen: + jump = 0 # don't need to jump yet + # positive as a potential jump start + last_seen[output_key] = permute_idx + else: + prev = last_seen[output_key] + if prev >= 0: # positive ==> it's a jump start + # jump to current idx, positive as the jump start + permutes[prev * permute_param + 5] = permute_idx + else: # it's already in a jump sequence, mark as negative + permutes[-prev * permute_param + 5] = -permute_idx + # mark last_seen negative since it's already in jump + last_seen[output_key] = -permute_idx + # it's a potential jump stop + jump = -total_permutes + + permutes[permute_idx * permute_param : permute_idx * permute_param + 6] = [ + input_tensor_idx, + output_tensor_idx, + input_start, + output_start, + length, + jump, + ] + permute_idx += 1 + output_start += length + out_lengths[output_tensor_idx] = output_start + + return ( + permutes, + in_lengths, + out_lengths, + ) + + def _values_string(values: torch.Tensor, start: int, end: int) -> str: size = values.size() if len(size) == 1: diff --git a/torchrec/sparse/tests/test_jagged_tensor.py b/torchrec/sparse/tests/test_jagged_tensor.py index 9efeb444c..5e49f3516 100644 --- a/torchrec/sparse/tests/test_jagged_tensor.py +++ b/torchrec/sparse/tests/test_jagged_tensor.py @@ -16,6 +16,7 @@ from torch.testing import FileCheck from torchrec.fx import symbolic_trace from torchrec.sparse.jagged_tensor import ( + _multi_remap_to_groups, _regroup_keyed_tensors, ComputeJTDictToKJT, ComputeKJTToJTDict, @@ -1374,6 +1375,170 @@ def test_permute_vb(self) -> None: ) self.assertEqual(permuted_jag_tensor.weights_or_none(), None) + def test_multi_remap_to_group(self) -> None: + keys = [["f1", "f2"], ["f3", "f4", "f5"], ["f6"]] + lengths = [[3, 4], [5, 6, 7], [8]] + groups = [["f1", "f3"], ["f2"], ["f4", "f1", "f6"], ["f1", "f5"]] + permutes, in_lengths, out_lengths = _multi_remap_to_groups( + keys, lengths, groups + ) + ref_permutes = [ + [0, 0, 0, 0, 3, 4], # f1, jump to 4, as a start + [1, 0, 0, 3, 5, 0], # f3 + [0, 1, 3, 0, 4, 0], # f2 + [1, 2, 5, 0, 6, 0], # f4 + [0, 2, 0, 6, 3, -6], # f1 jump to 6, as in a jump sequence + [2, 2, 0, 9, 8, 0], # f6 + [0, 3, 0, 0, 3, -8], # f1 jump stop, as out of boundary + [1, 3, 11, 3, 7, 0], # f5 + ] + self.assertEqual(permutes, [i for p in ref_permutes for i in p]) + self.assertEqual(in_lengths, [7, 18, 8]) + self.assertEqual(out_lengths, [8, 4, 17, 10]) + + def test_multi_permute_forward_cpu(self) -> None: + batch_size = 5 + keys = [["f1", "f2"], ["f3", "f4", "f5"], ["f6"]] + lengths = [[3, 4], [5, 6, 7], [8]] + groups = [["f1", "f3"], ["f2"], ["f4", "f1", "f6"], ["f1", "f5"]] + values = [ + torch.randn(batch_size, sum(lens), device="cpu", requires_grad=True) + for lens in lengths + ] + permutes, in_lengths, out_lengths = _multi_remap_to_groups( + keys, lengths, groups + ) + refs = [[] for _ in groups] + for i in range(len(permutes) // 6): + in_idx, out_idx, in_start, _, length, _ = permutes[i * 6 : i * 6 + 6] + refs[out_idx].append(values[in_idx][:, in_start : (in_start + length)]) + refs = [torch.cat(ref, dim=1) for ref in refs] + outputs = torch.ops.fbgemm.permute_multi_embedding( + values, permutes, in_lengths, out_lengths + ) + for out, ref in zip(outputs, refs): + self.assertTrue(torch.allclose(out, ref)) + + def test_multi_permute_forward_meta(self) -> None: + batch_size = 5 + keys = [["f1", "f2"], ["f3", "f4", "f5"], ["f6"]] + lengths = [[3, 4], [5, 6, 7], [8]] + groups = [["f1", "f3"], ["f2"], ["f4", "f1", "f6"], ["f1", "f5"]] + values = [ + torch.randn(batch_size, sum(lens), device="meta", requires_grad=True) + for lens in lengths + ] + permutes, in_lengths, out_lengths = _multi_remap_to_groups( + keys, lengths, groups + ) + refs = [[] for _ in groups] + for i in range(len(permutes) // 6): + in_idx, out_idx, in_start, _, length, _ = permutes[i * 6 : i * 6 + 6] + refs[out_idx].append(values[in_idx][:, in_start : (in_start + length)]) + refs = [torch.cat(ref, dim=1) for ref in refs] + outputs = torch.ops.fbgemm.permute_multi_embedding( + values, permutes, in_lengths, out_lengths + ) + for out, ref in zip(outputs, refs): + self.assertEqual(out.shape, ref.shape) + + def test_multi_permute_forward_gpu(self) -> None: + batch_size = 5 + keys = [["f1", "f2"], ["f3", "f4", "f5"], ["f6"]] + lengths = [[3, 4], [5, 6, 7], [8]] + groups = [["f1", "f3"], ["f2"], ["f4", "f1", "f6"], ["f1", "f5"]] + values = [ + torch.randn(batch_size, sum(lens), device="cuda", requires_grad=True) + for lens in lengths + ] + permutes, in_lengths, out_lengths = _multi_remap_to_groups( + keys, lengths, groups + ) + refs = [[] for _ in groups] + for i in range(len(permutes) // 6): + in_idx, out_idx, in_start, _, length, _ = permutes[i * 6 : i * 6 + 6] + refs[out_idx].append(values[in_idx][:, in_start : (in_start + length)]) + refs = [torch.cat(ref, dim=1) for ref in refs] + outputs = torch.ops.fbgemm.permute_multi_embedding( + values, permutes, in_lengths, out_lengths + ) + for out, ref in zip(outputs, refs): + self.assertTrue(torch.allclose(out, ref)) + + def test_multi_permute_backward_cpu(self) -> None: + batch_size = 5 + keys = [["f1", "f2"], ["f3", "f4", "f5"], ["f6"]] + lengths = [[3, 4], [5, 6, 7], [8]] + groups = [["f1", "f3"], ["f2"], ["f4", "f1", "f6"], ["f1", "f5"]] + values = [ + torch.randn(batch_size, sum(lens), device="cpu", requires_grad=True) + for lens in lengths + ] + ref_values = [v.detach() for v in values] + for v in ref_values: + v.requires_grad = True + permutes, in_lengths, out_lengths = _multi_remap_to_groups( + keys, lengths, groups + ) + refs = [[] for _ in groups] + for i in range(len(permutes) // 6): + in_idx, out_idx, in_start, _, length, _ = permutes[i * 6 : i * 6 + 6] + refs[out_idx].append(ref_values[in_idx][:, in_start : (in_start + length)]) + refs = [torch.cat(ref, dim=1) for ref in refs] + outputs = torch.ops.fbgemm.permute_multi_embedding( + values, permutes, in_lengths, out_lengths + ) + for out, ref in zip(outputs, refs): + self.assertTrue(torch.allclose(out, ref)) + + ref_loss, loss = refs[0].sum(), outputs[0].sum() + for i in range(1, len(refs)): + ref_loss += (i + 1.1) * refs[i].sum() + loss += (i + 1.1) * outputs[i].sum() + ref_loss.backward() + loss.backward() + for val, ref in zip(values, ref_values): + val_grad, ref_grad = val.grad, ref.grad + assert isinstance(val_grad, torch.Tensor) + self.assertTrue(torch.allclose(val_grad, ref_grad)) + + def test_multi_permute_backward_gpu(self) -> None: + batch_size = 2048 + keys = [["f1", "f2"], ["f3", "f4", "f5"], ["f6"]] + lengths = [[96, 256], [512, 128, 768], [1024]] + groups = [["f1", "f3"], ["f2"], ["f4", "f1", "f6"], ["f1", "f5"]] + values = [ + torch.randn(batch_size, sum(lens), device="cuda", requires_grad=True) + for lens in lengths + ] + ref_values = [v.detach() for v in values] + for v in ref_values: + v.requires_grad = True + permutes, in_lengths, out_lengths = _multi_remap_to_groups( + keys, lengths, groups + ) + refs = [[] for _ in groups] + for i in range(len(permutes) // 6): + in_idx, out_idx, in_start, _, length, _ = permutes[i * 6 : i * 6 + 6] + refs[out_idx].append(ref_values[in_idx][:, in_start : (in_start + length)]) + refs = [torch.cat(ref, dim=1) for ref in refs] + outputs = torch.ops.fbgemm.permute_multi_embedding( + values, permutes, in_lengths, out_lengths + ) + for out, ref in zip(outputs, refs): + self.assertTrue(torch.allclose(out, ref)) + + ref_loss, loss = refs[0].sum(), outputs[0].sum() + for i in range(1, len(refs)): + ref_loss += (i + 1.1) * refs[i].sum() + loss += (i + 1.1) * outputs[i].sum() + ref_loss.backward() + loss.backward() + for val, ref in zip(values, ref_values): + val_grad, ref_grad = val.grad, ref.grad + assert isinstance(val_grad, torch.Tensor) + self.assertTrue(torch.allclose(val_grad, ref_grad)) + def test_permute_duplicates(self) -> None: values = torch.Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]) lengths = torch.IntTensor([0, 2, 0, 1, 1, 1, 0, 3, 0]) @@ -1650,8 +1815,6 @@ def test_string_vb(self) -> None: stride_per_key_per_rank=stride_per_key_per_rank, ) - print(str(jag_tensor)) - self.assertEqual( str(jag_tensor), """\