Enable Pyrefly for deeplearning/fbgemm (#6099)

Summary:
Pull Request resolved: #6099

X-link: https://github.com/facebookresearch/FBGEMM/pull/3003

Enables Pyrefly for `deeplearning/fbgemm` by setting `python.set_pyrefly(True)` in the subdirectory PACKAGE file and adding targeted `# pyrefly: ignore` annotations. Part of splitting D113608658 into a per-subdirectory stack; the final diff performs the directory-level flip.

#buildall

Reviewed By: maggiemoss

Differential Revision: D113962589

fbshipit-source-id: aea111f522e6951bdac4abc0d5d28484326f6cb7

Author: QuantumManiac

fbgemm_gpu/bench/jagged_tensor_benchmark.py

@@ -1089,6 +1089,7 @@ def keyed_jagged_index_select_dim1_ref(
         )
     if baseline:
         time_ref, _ = benchmark_torch_function(
+            # pyrefly: ignore [unbound-name]
             functools.partial(output_ref.backward, retain_graph=True),
             (grad,),
             iters=iters,

fbgemm_gpu/bench/merge_embeddings_benchmark.py

@@ -94,6 +94,7 @@ def generate_requests(
         rs.append(
             get_table_batched_offsets_from_dense(all_indices.view(T, B, L), gpu_num)
         )
+    # pyrefly: ignore [bad-return]
     return rs
 
 

fbgemm_gpu/bench/sparse_ops_benchmark.py

@@ -1402,6 +1402,7 @@ def batch_group_index_select_bwd(
         )
 
     if timeline:
+        # pyrefly: ignore [missing-attribute]
         prof.export_chrome_trace("index_select_fwd_trace.json")
 
     grads = [torch.rand_like(out) for out in out_pyt]
@@ -1442,6 +1443,7 @@ def batch_group_index_select_bwd(
         )
 
     if timeline:
+        # pyrefly: ignore [missing-attribute]
         prof.export_chrome_trace("index_select_bwd_trace.json")
 
     logging.info(
@@ -1722,10 +1724,13 @@ def permute_1d_sparse_data_bench(
     embedding tables.
     """
     if index_dtype == "int":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.int32
     elif index_dtype == "int64":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.int64
     elif index_dtype == "float":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.float32
     else:
         raise RuntimeError(f"Does not support data type {index_dtype}")
@@ -1744,8 +1749,10 @@ def permute_1d_sparse_data_bench(
     total_indices = int(lengths.sum().item())
     # Generate indices
     if index_dtype == torch.float32:
+        # pyrefly: ignore [no-matching-overload]
         indices = torch.rand(total_indices, dtype=index_dtype, device=device)
     else:
+        # pyrefly: ignore [no-matching-overload]
         indices = torch.randint(
             low=0,
             high=2**31 - 1,
@@ -1881,14 +1888,19 @@ def permute_2d_sparse_data_bench(
     systems to reorder embedding tables.
     """
     if index_dtype == "int":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.int32
     elif index_dtype == "int64":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.int64
     elif index_dtype == "float":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.float32
     elif index_dtype == "bf16":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.bfloat16
     elif index_dtype == "fp16":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.float16
     else:
         raise RuntimeError(f"Does not support data type {index_dtype}")
@@ -1909,9 +1921,12 @@ def permute_2d_sparse_data_bench(
         * emb_dim
     )
     total_indices = int(lengths.sum().item())
+    # pyrefly: ignore [missing-attribute]
     if index_dtype.is_floating_point:
+        # pyrefly: ignore [no-matching-overload]
         indices = torch.rand(total_indices, dtype=index_dtype, device=device)
     else:
+        # pyrefly: ignore [no-matching-overload]
         indices = torch.randint(
             low=0,
             high=2**31 - 1,

fbgemm_gpu/bench/tbe/split_table_batched_embeddings_benchmark.py

@@ -2047,6 +2047,7 @@ def device_from_files(  # noqa C901
     _Ds: list[int]
     emb_loc: list[EmbeddingLocation]
     compute_device: list[ComputeDevice]
+    # pyrefly: ignore [bad-assignment]
     _Es, _Ds, emb_loc, compute_device = zip(*embedding_specs)
 
     # Determine location suffix for trace URL based on embedding locations

fbgemm_gpu/bench/tbe/tbe_cache_benchmark.py

@@ -272,12 +272,14 @@ def lxu_cache_lookup(
     tbe: nn.Module = IntNBitTableBatchedEmbeddingBagsCodegen(
         embedding_specs, cache_load_factor=cache_load_factor
     )
+    # pyrefly: ignore [not-callable]
     tbe.fill_random_weights()
 
     # Imitate execution flow by performing prefetching once.
     indices, offsets = create_request(
         num_tables, num_embeddings, batch, avg_pooling_factor
     )
+    # pyrefly: ignore [not-callable]
     tbe.prefetch(indices, offsets)
 
     linearized_indices = torch.ops.fbgemm.linearize_cache_indices(
@@ -333,6 +335,7 @@ def lru_cache_populate_byte(
     cc: nn.Module = IntNBitTableBatchedEmbeddingBagsCodegen(
         embedding_specs, cache_load_factor=cache_load_factor
     )
+    # pyrefly: ignore [not-callable]
     cc.fill_random_weights()
 
     warm_up_requests = []
@@ -389,6 +392,7 @@ def populate(linear_indices: Tensor) -> None:
     replay_cc: nn.Module = IntNBitTableBatchedEmbeddingBagsCodegen(
         embedding_specs, cache_load_factor=cache_load_factor
     )
+    # pyrefly: ignore [not-callable]
     replay_cc.fill_random_weights()
 
     replay_timestep: int = 1
@@ -464,6 +468,7 @@ def lfu_cache_populate_byte(
         cache_load_factor=cache_load_factor,
         cache_algorithm=CacheAlgorithm.LFU,
     )
+    # pyrefly: ignore [not-callable]
     cc.fill_random_weights()
 
     warm_up_requests = []
@@ -517,6 +522,7 @@ def populate(linear_indices: Tensor) -> None:
         cache_load_factor=cache_load_factor,
         cache_algorithm=CacheAlgorithm.LFU,
     )
+    # pyrefly: ignore [not-callable]
     replay_cc.fill_random_weights()
 
     def replay_populate(linear_indices: Tensor) -> None:

fbgemm_gpu/bench/tbe/tbe_inference_benchmark.py

@@ -874,6 +874,7 @@ def nbit_device_with_spec(  # noqa C901
                 TBERequest(
                     req.indices.cpu().int(),
                     req.offsets.cpu().int(),
+                    # pyrefly: ignore [missing-attribute]
                     req.per_sample_weigths.cpu() if req.per_sample_weights else None,
                 )
                 for req in requests
@@ -1197,6 +1198,7 @@ def nbit_uvm(
     )
 
     if T_gpu > 0:
+        # pyrefly: ignore [unbound-name]
         nparams_byte = sum(w.numel() for (w, _) in emb_mixed.split_embedding_weights())
         logging.info(
             f"{weights_precision} Embedding tables: {E * T_gpu + E_uvm * T_uvm} rows, {nparams_byte / param_size_multiplier / 1.0e9: .2f} GParam, "
@@ -1647,6 +1649,7 @@ def nbit_cache(  # noqa C901
             )
             for d in Ds
         ],
+        # pyrefly: ignore [not-callable]
         record_cache_metrics=RecordCacheMetrics(
             record_cache_miss_counter, record_tablewise_cache_miss
         ),

fbgemm_gpu/bench/tbe/tbe_training_benchmark.py

@@ -639,6 +639,7 @@ def device_with_speclist(  # noqa C901
     # pyre-ignore[53]
     def _kineto_trace_handler(p: profile, phase: str) -> None:
         p.export_chrome_trace(
+            # pyrefly: ignore [missing-attribute]
             benchconfig.trace_url.format(
                 emb_op_type=emb_op_type, phase=phase, ospid=os.getpid()
             )

fbgemm_gpu/codegen/genscript/generate_backward_split.py

@@ -26,6 +26,7 @@
     from .optimizer_args import annotation_dict, OptimizerArgsSet
     from .scripts_argsparse import args
 except ImportError:
+    # pyrefly: ignore [missing-import]
     from optimizers import *
 
     # pyre-ignore[21]
@@ -118,6 +119,7 @@ def generate_backward_split_gpu(**kwargs: Any) -> None:
         ]:
             BackwardSplitGenerator.render_backward_templates(
                 template_filepath,
+                # pyrefly: ignore [bad-argument-type]
                 optimizer,
                 filename_format,
                 kwargs,
@@ -141,6 +143,7 @@ def generate_backward_split_gpu(**kwargs: Any) -> None:
             ]:
                 BackwardSplitGenerator.render_backward_templates(
                     template_filepath,
+                    # pyrefly: ignore [bad-argument-type]
                     optimizer,
                     filename_format,
                     kwargs,
@@ -241,7 +244,7 @@ def generate_backward_split_cpu(**kwargs: Any) -> None:
         if kwargs.get("has_cpu_support"):
             CodeTemplate.load(
                 "training/backward/embedding_backward_split_cpu_approx_template.cpp"
-                if "approx" in optimizer
+                if "approx" in optimizer  # pyrefly: ignore [not-iterable]
                 else "training/backward/embedding_backward_split_cpu_template.cpp"
             ).write(f"gen_embedding_backward_{optimizer}_split_cpu.cpp", **kwargs)
 

fbgemm_gpu/codegen/genscript/generate_embedding_optimizer.py

@@ -36,6 +36,7 @@ def generate_embedding_optimizer(**kwargs: Any) -> None:
 
         optimizer = kwargs.get("optimizer")
         kwargs["optimizer_class_name"] = "".join(
+            # pyrefly: ignore [missing-attribute]
             [optim.capitalize() for optim in optimizer.split("_")]
         )
         kwargs["args"] = kwargs["args"].cuda

fbgemm_gpu/codegen/genscript/jinja_environment.py

@@ -58,20 +58,26 @@
 # BT_block_size * 4 * 4 * 32 * (max_D // 128) <= 64 * 1024 (V100) or 96 * 1024 (A100)
 # Since BT_block_size >= 1, max_D <= 16K (V100) or 24K (A100).
 # Note that if we increase max_D, it will increase the compilation time significantly.
+# pyrefly: ignore [unsupported-operation]
 env.globals["max_embedding_dim"] = 2048
 
 # Max embedding dimension for legacy embedding kernels. TBE v2 can support
 # larger max embedding dimension.
+# pyrefly: ignore [unsupported-operation]
 env.globals["legacy_max_embedding_dim"] = 1024
 
 # An optimization for ROCm
+# pyrefly: ignore [unsupported-operation]
 env.globals["items_per_warp"] = 128 if args.is_rocm is False else 256
 
 # The fixed max vectors per thread for different kernels.  The numbers were
 # derived from empirical studies
+# pyrefly: ignore [unsupported-operation]
 env.globals["fixed_max_vecs_per_thread"] = {"backward": 2, "backward_indice_weights": 6}
 
+# pyrefly: ignore [unsupported-operation]
 env.globals["dense"] = False
+# pyrefly: ignore [unsupported-operation]
 env.globals["is_rocm"] = args.is_rocm
 
 
@@ -343,15 +349,24 @@ def compute_global_weight_decay(is_global_weight_decay_kernel: bool) -> str:
 ################################################################################
 
 env.globals["generate_optimized_grad_sum_loop_access"] = (
+    # pyrefly: ignore [unsupported-operation]
     generate_optimized_grad_sum_loop_access
 )
+# pyrefly: ignore [unsupported-operation]
 env.globals["get_max_vecs_template_configs"] = get_max_vecs_template_configs
+# pyrefly: ignore [unsupported-operation]
 env.globals["dispatch_optimal_kernel"] = dispatch_optimal_kernel
+# pyrefly: ignore [unsupported-operation]
 env.globals["dispatch_non_vec_blocking_kernel"] = dispatch_non_vec_blocking_kernel
+# pyrefly: ignore [unsupported-operation]
 env.globals["dispatch_vec_blocking_kernel"] = dispatch_vec_blocking_kernel
+# pyrefly: ignore [unsupported-operation]
 env.globals["is_valid_forward_config"] = is_valid_forward_config
+# pyrefly: ignore [unsupported-operation]
 env.globals["has_experimental_support"] = has_experimental_support
+# pyrefly: ignore [unsupported-operation]
 env.globals["is_valid_gwd_config"] = is_valid_gwd_config
+# pyrefly: ignore [unsupported-operation]
 env.globals["compute_global_weight_decay"] = compute_global_weight_decay
 
 ################################################################################