[TP] Add API logging for TP high level API

[ghstack-poisoned]
pytorch · May 24, 2023 · e0921a2 · e0921a2
1 parent e3d97b6
commit e0921a2
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/torch/distributed/tensor/parallel/api.py b/torch/distributed/tensor/parallel/api.py
@@ -80,6 +80,8 @@ def parallelize_module(  # type: ignore[return]
         granularity, you need to pass in a dict of module FQN and parallel style instead.
     """
 
+    torch._C._log_api_usage_once("torch.distributed.tensor.parallel.parallelize_module")
+
     if device_mesh.ndim > 1:
         device_mesh = _create_1d_device_mesh(device_mesh, tp_mesh_dim)
 

diff --git a/torch/distributed/tensor/parallel/fsdp.py b/torch/distributed/tensor/parallel/fsdp.py
@@ -43,6 +43,9 @@ def enable_2d_with_fsdp() -> bool:
     Return:
         A `bool` indicated whether extension registration succeeds or not.
     """
+
+    torch._C._log_api_usage_once("torch.distributed.tensor.parallel.enable_2d_with_fsdp")
+
     try:
         from torch.distributed.fsdp._fsdp_extensions import (
             _set_fsdp_extensions,