fbgemm-gpu-nightly-cpu 2025.3.27__cp311-cp311-manylinux_2_28_aarch64.whl → 2026.1.29__cp311-cp311-manylinux_2_28_aarch64.whl

fbgemm_gpu/split_embedding_codegen_lookup_invokers/lookup_sgd.py

@@ -4,6 +4,8 @@
 ## Template Source: training/python/split_embedding_codegen_lookup_invoker.template
 ################################################################################
 
+__template_source_file__ = "training/python/split_embedding_codegen_lookup_invoker.template"
+
 #!/usr/bin/env python3
 
 # Copyright (c) Meta Platforms, Inc. and affiliates.
@@ -15,6 +17,12 @@
 
 import torch
 from .lookup_args import *
+
+
+
+
+
+
 def invoke(
     common_args: CommonArgs,
     optimizer_args: OptimizerArgs,
@@ -23,84 +31,177 @@ def invoke(
     # only pass prev_iter_dev since prev_iter is never created on UVM
     prev_iter_dev: Optional[torch.Tensor] = None,
     gwd_lower_bound: float = 0.0,
+    mixed_D: bool = True,
 ) -> torch.Tensor:
-
+    # host_weights is only used for CPU training
+    use_cpu = common_args.host_weights.numel() > 0
     vbe_metadata = common_args.vbe_metadata
-    if (common_args.host_weights.numel() > 0):
-        T = common_args.D_offsets.numel() - 1
-        vbe: bool = vbe_metadata.B_offsets is not None
-        if vbe:
-            # create offsets with fixed batch size max_B
-            # not efficient but for now we just need a functional implementation for CPU
-            max_B = vbe_metadata.max_B
-            offsets = torch.empty([T * max_B + 1], dtype=common_args.offsets.dtype, device=common_args.offsets.device)
-            for t in range(T):
-                B_offsets = vbe_metadata.B_offsets
-                assert isinstance(B_offsets, torch.Tensor)
-                begin = B_offsets[t]
-                end = B_offsets[t + 1]
-                offsets[t * max_B : t * max_B + end - begin] = common_args.offsets[begin : end]
-                offsets[t * max_B + end - begin : (t + 1) * max_B] = common_args.offsets[end]
-            offsets[-1] = common_args.offsets[-1]
-        else:
-            offsets = common_args.offsets
-        output = torch.ops.fbgemm.split_embedding_codegen_lookup_sgd_function_cpu(
-            # common_args
-            host_weights=common_args.host_weights,
-            weights_placements=common_args.weights_placements,
-            weights_offsets=common_args.weights_offsets,
-            D_offsets=common_args.D_offsets,
-            total_D=common_args.total_D,
-            max_D=common_args.max_D,
-            hash_size_cumsum=common_args.hash_size_cumsum,
-            total_hash_size_bits=common_args.total_hash_size_bits,
-            indices=common_args.indices,
-            offsets=offsets,
-            pooling_mode=common_args.pooling_mode,
-            indice_weights=common_args.indice_weights,
-            feature_requires_grad=common_args.feature_requires_grad,
-            # optimizer_args
-            gradient_clipping = optimizer_args.gradient_clipping,
-            max_gradient=optimizer_args.max_gradient,
-            stochastic_rounding=optimizer_args.stochastic_rounding,
-            learning_rate=optimizer_args.learning_rate,
-            # momentum1
-            # momentum2
-            # prev_iter
-            # row_counter
-            # iter
-            # max counter
-        )
-        if vbe:
-            output_new = torch.empty([vbe_metadata.output_size], dtype=output.dtype, device=output.device)
-            B_offsets_rank_per_feature = vbe_metadata.B_offsets_rank_per_feature
-            assert isinstance(B_offsets_rank_per_feature, torch.Tensor)
-            output_offsets_feature_rank = vbe_metadata.output_offsets_feature_rank
-            assert isinstance(output_offsets_feature_rank, torch.Tensor)
-            R = B_offsets_rank_per_feature.size(1) - 1
-            for r in range(R):
-                D_offset = 0
-                for t in range(T):
-                    o_begin = output_offsets_feature_rank[r * T + t].item()
-                    o_end = output_offsets_feature_rank[r * T + t + 1].item()
-                    D = common_args.D_offsets[t + 1].item() - common_args.D_offsets[t].item()
-                    b_begin = B_offsets_rank_per_feature[t][r].item()
-                    b_end = B_offsets_rank_per_feature[t][r + 1].item()
-                    assert o_end - o_begin == (b_end - b_begin) * D
-                    output_new[o_begin : o_end] = output[b_begin : b_end, D_offset : D_offset + D].flatten()
-                    D_offset += D
-            return output_new
-        else:
-            return output
-
-    return torch.ops.fbgemm.split_embedding_codegen_lookup_sgd_function(
+
+    # pack weights
+    weights = [
+        common_args.dev_weights,
+        common_args.uvm_weights,
+        common_args.weights_placements,
+        common_args.weights_offsets,
+        common_args.lxu_cache_weights,
+    ] if not use_cpu else [
+        common_args.host_weights,
+        common_args.weights_placements,
+        common_args.weights_offsets,
+    ]
+    dict_aux_tensor: Dict[str, Optional[torch.Tensor]] = {
+        "B_offsets": vbe_metadata.B_offsets,
+        "vbe_output_offsets_feature_rank": vbe_metadata.output_offsets_feature_rank,
+        "vbe_B_offsets_rank_per_feature": vbe_metadata.B_offsets_rank_per_feature,
+        "lxu_cache_locations": common_args.lxu_cache_locations,
+        "uvm_cache_stats": common_args.uvm_cache_stats,
+        "vbe_output_offsets" : vbe_metadata.vbe_output_offsets,
+    }
+
+    dict_aux_int: Dict[str, int] = {
+        "iter": iter,
+        "info_B_num_bits": common_args.info_B_num_bits,
+        "info_B_mask": common_args.info_B_mask,
+    }
+
+    dict_aux_float: Dict[str, float] = {
+        "gwd_lower_bound": gwd_lower_bound,
+    }
+
+    dict_aux_bool: Dict[str, bool] = {
+        "is_experimental_tbe": common_args.is_experimental,
+        "use_uniq_cache_locations_bwd": common_args.use_uniq_cache_locations_bwd,
+        "use_homogeneous_placements": common_args.use_homogeneous_placements,
+        "apply_global_weight_decay": apply_global_weight_decay,
+        "mixed_D": mixed_D,
+    }
+    dict_optim_int: Dict[str, int] = {}
+    dict_optim_float: Dict[str, float] = {}
+    dict_optim_bool: Dict[str, bool] = {}
+
+    # Explicitly pass only prev_iter_dev for global weight decay, unless it already exists in optim arg
+    dict_aux_tensor["prev_iter_dev"] = prev_iter_dev
+
+
+    # optimizer_args # if optimizer == none
+    dict_aux_bool["gradient_clipping"] = optimizer_args.gradient_clipping
+    dict_aux_float["max_gradient"] = optimizer_args.max_gradient
+    dict_aux_bool["stochastic_rounding"] = optimizer_args.stochastic_rounding
+
+    
+    aux_tensor: List[Optional[torch.Tensor]] = []
+    assert "B_offsets" in dict_aux_tensor, (
+        "B_offsets must be in dict_aux_tensor. "
+        "Please check the frontend and backend version. "
+    )
+    aux_tensor.append(dict_aux_tensor["B_offsets"])
+    assert "vbe_output_offsets_feature_rank" in dict_aux_tensor, (
+        "vbe_output_offsets_feature_rank must be in dict_aux_tensor. "
+        "Please check the frontend and backend version. "
+    )
+    aux_tensor.append(dict_aux_tensor["vbe_output_offsets_feature_rank"])
+    assert "vbe_B_offsets_rank_per_feature" in dict_aux_tensor, (
+        "vbe_B_offsets_rank_per_feature must be in dict_aux_tensor. "
+        "Please check the frontend and backend version. "
+    )
+    aux_tensor.append(dict_aux_tensor["vbe_B_offsets_rank_per_feature"])
+    assert "lxu_cache_locations" in dict_aux_tensor, (
+        "lxu_cache_locations must be in dict_aux_tensor. "
+        "Please check the frontend and backend version. "
+    )
+    aux_tensor.append(dict_aux_tensor["lxu_cache_locations"])
+    assert "uvm_cache_stats" in dict_aux_tensor, (
+        "uvm_cache_stats must be in dict_aux_tensor. "
+        "Please check the frontend and backend version. "
+    )
+    aux_tensor.append(dict_aux_tensor["uvm_cache_stats"])
+    assert "prev_iter_dev" in dict_aux_tensor, (
+        "prev_iter_dev must be in dict_aux_tensor. "
+        "Please check the frontend and backend version. "
+    )
+    aux_tensor.append(dict_aux_tensor["prev_iter_dev"])
+    assert "vbe_output_offsets" in dict_aux_tensor, (
+        "vbe_output_offsets must be in dict_aux_tensor. "
+        "Please check the frontend and backend version. "
+    )
+    aux_tensor.append(dict_aux_tensor["vbe_output_offsets"])
+    
+    aux_int: List[int] = []
+    assert "iter" in dict_aux_int, (
+        "iter must be in dict_aux_int. "
+        "Please check the frontend and backend version. "
+    )
+    aux_int.append(dict_aux_int["iter"])
+    assert "info_B_num_bits" in dict_aux_int, (
+        "info_B_num_bits must be in dict_aux_int. "
+        "Please check the frontend and backend version. "
+    )
+    aux_int.append(dict_aux_int["info_B_num_bits"])
+    assert "info_B_mask" in dict_aux_int, (
+        "info_B_mask must be in dict_aux_int. "
+        "Please check the frontend and backend version. "
+    )
+    aux_int.append(dict_aux_int["info_B_mask"])
+    
+    aux_float: List[float] = []
+    assert "gwd_lower_bound" in dict_aux_float, (
+        "gwd_lower_bound must be in dict_aux_float. "
+        "Please check the frontend and backend version. "
+    )
+    aux_float.append(dict_aux_float["gwd_lower_bound"])
+    assert "max_gradient" in dict_aux_float, (
+        "max_gradient must be in dict_aux_float. "
+        "Please check the frontend and backend version. "
+    )
+    aux_float.append(dict_aux_float["max_gradient"])
+    
+    aux_bool: List[bool] = []
+    assert "is_experimental_tbe" in dict_aux_bool, (
+        "is_experimental_tbe must be in dict_aux_bool. "
+        "Please check the frontend and backend version. "
+    )
+    aux_bool.append(dict_aux_bool["is_experimental_tbe"])
+    assert "use_uniq_cache_locations_bwd" in dict_aux_bool, (
+        "use_uniq_cache_locations_bwd must be in dict_aux_bool. "
+        "Please check the frontend and backend version. "
+    )
+    aux_bool.append(dict_aux_bool["use_uniq_cache_locations_bwd"])
+    assert "use_homogeneous_placements" in dict_aux_bool, (
+        "use_homogeneous_placements must be in dict_aux_bool. "
+        "Please check the frontend and backend version. "
+    )
+    aux_bool.append(dict_aux_bool["use_homogeneous_placements"])
+    assert "apply_global_weight_decay" in dict_aux_bool, (
+        "apply_global_weight_decay must be in dict_aux_bool. "
+        "Please check the frontend and backend version. "
+    )
+    aux_bool.append(dict_aux_bool["apply_global_weight_decay"])
+    assert "gradient_clipping" in dict_aux_bool, (
+        "gradient_clipping must be in dict_aux_bool. "
+        "Please check the frontend and backend version. "
+    )
+    aux_bool.append(dict_aux_bool["gradient_clipping"])
+    assert "stochastic_rounding" in dict_aux_bool, (
+        "stochastic_rounding must be in dict_aux_bool. "
+        "Please check the frontend and backend version. "
+    )
+    aux_bool.append(dict_aux_bool["stochastic_rounding"])
+    assert "mixed_D" in dict_aux_bool, (
+        "mixed_D must be in dict_aux_bool. "
+        "Please check the frontend and backend version. "
+    )
+    aux_bool.append(dict_aux_bool["mixed_D"])
+
+    # optim_int
+    # optim_float
+    # ['learning_rate_tensor']
+    # optim_bool
+
+    return torch.ops.fbgemm.split_embedding_codegen_lookup_sgd_function_pt2(
         # common_args
         placeholder_autograd_tensor=common_args.placeholder_autograd_tensor,
-        dev_weights=common_args.dev_weights,
-        uvm_weights=common_args.uvm_weights,
-        lxu_cache_weights=common_args.lxu_cache_weights,
-        weights_placements=common_args.weights_placements,
-        weights_offsets=common_args.weights_offsets,
+        # weights
+        weights=weights,
         D_offsets=common_args.D_offsets,
         total_D=common_args.total_D,
         max_D=common_args.max_D,
@@ -111,35 +212,30 @@ def invoke(
         pooling_mode=common_args.pooling_mode,
         indice_weights=common_args.indice_weights,
         feature_requires_grad=common_args.feature_requires_grad,
-        lxu_cache_locations=common_args.lxu_cache_locations,
-        uvm_cache_stats=common_args.uvm_cache_stats,
+        output_dtype=common_args.output_dtype,
         # VBE metadata
-        B_offsets=vbe_metadata.B_offsets,
-        vbe_output_offsets_feature_rank=vbe_metadata.output_offsets_feature_rank,
-        vbe_B_offsets_rank_per_feature=vbe_metadata.B_offsets_rank_per_feature,
         max_B=vbe_metadata.max_B,
         max_B_feature_rank=vbe_metadata.max_B_feature_rank,
         vbe_output_size=vbe_metadata.output_size,
-        # optimizer_args
-        gradient_clipping = optimizer_args.gradient_clipping,
-        max_gradient=optimizer_args.max_gradient,
-        stochastic_rounding=optimizer_args.stochastic_rounding, # if optimizer == none
-        # V1 interface still accepts learning_rate as float
-        learning_rate=optimizer_args.learning_rate,
+        vbe_output=vbe_metadata.vbe_output,
+        # aux_tensor
+        aux_tensor=aux_tensor,
+        # aux_int
+        aux_int=aux_int,
+        # aux_float
+        aux_float=aux_float,
+        # aux_bool
+        aux_bool=aux_bool,
+        learning_rate_tensor=common_args.learning_rate_tensor,
+
         # momentum1
         # momentum2
         # prev_iter
-        
-        prev_iter_dev=prev_iter_dev,
         # row_counter
-        # iter
-        iter=iter,
-        # max counter
+        # optim_tensor
+        # optim_int
+        # optim_float
+        # optim_bool
+        # optim symint args
         # total_unique_indices
-        output_dtype=common_args.output_dtype,
-        is_experimental=common_args.is_experimental,
-        use_uniq_cache_locations_bwd=common_args.use_uniq_cache_locations_bwd,
-        use_homogeneous_placements=common_args.use_homogeneous_placements,
-        apply_global_weight_decay=apply_global_weight_decay,
-        gwd_lower_bound=gwd_lower_bound,
     )