Log forward intermediates hashes w/pp vs w/o pp

stack-info: PR: #246, branch: xmfan/stack/20

Author: xmfan

autoparallel/api.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import copy
+import functools
 import itertools
 import warnings
 from contextlib import ExitStack, contextmanager
@@ -42,7 +43,11 @@
 )
 from .init_weights import hook_params_setters
 from .optimize_sharding import ShardingOptimizer
-from .utils import _get_device_from_mesh
+from .utils import (
+    NumericsLogger,
+    _get_device_from_mesh,
+    debug_boxed_nop_preserve_node_meta,
+)
 
 _APPLY_VIEW_MM_VIEW_PATTERN = False
 
@@ -193,6 +198,7 @@ def __init__(
         ac_stage_size_in_GiB: Optional[Union[float, str]] = "auto",
         reshard_after_forward: bool = True,
         dynamic: bool = False,
+        numerics_logger: NumericsLogger | None = None,
         **kwargs,
     ):
         self.stack = ExitStack()
@@ -220,7 +226,14 @@ def __init__(
         self.model = move_to_fake(model, self.fake_mode, device)
         self.input_fn = input_fn
         self.mesh = mesh
-        self.compiler_fn = compile_fx_inner if compile else boxed_nop_preserve_node_meta
+        if compile:
+            self.compiler_fn = compile_fx_inner
+        elif numerics_logger:
+            self.compiler_fn = functools.partial(
+                debug_boxed_nop_preserve_node_meta, numerics_logger=numerics_logger
+            )
+        else:
+            self.compiler_fn = boxed_nop_preserve_node_meta
         self.enable_ac = enable_ac
         self.ac_stage_size_in_GiB = ac_stage_size_in_GiB
         self.reshard_after_forward = reshard_after_forward

autoparallel/utils.py

@@ -341,7 +341,7 @@ def log(self, node: str, args: Iterable[Any], inputs_or_outputs: str):
                 continue
 
             self._logs.append(
-                f"{node=}, {inputs_or_outputs}[{i}]={torch.hash_tensor(arg)}"
+                f"{node=}, {inputs_or_outputs}[{i}]={torch.hash_tensor(arg)} nan={torch.any(torch.isnan(arg))}"
             )
 
     def run_node(self, n: torch.fx.Node) -> Any:
@@ -429,6 +429,18 @@ def log_model_weights(self, parallel_mod):
 
             print(f"Weight hashes written to {path}")
 
+    def log_fw_intermediates(self, logs):
+        rank = torch.distributed.get_rank()
+        path = self.dir / f"rank_{rank}_fw_intermediates.log"
+        with open(path, "a") as f:
+            f.write("\n".join(logs) + "\n")
+
+    def log_forward_output(self, fw_out):
+        if self.rank == 0:
+            path = self.dir / "fw_out.log"
+            with open(path, "a") as f:
+                f.write(f"fw_out={hash_tensor(fw_out)}\n")
+
     def log_pp_model_weights(self, orig_mod, stage_mods, num_world_stages, ranks):
         path = self.dir / "pp_weights.log"
 
@@ -463,3 +475,17 @@ def log_pp_model_weights(self, orig_mod, stage_mods, num_world_stages, ranks):
 
         if self.rank == 0:
             print(f"Weight hashes written to {path}")
+
+
+def debug_boxed_nop_preserve_node_meta(fx_g, example_inputs, numerics_logger):
+    def run(args):
+        with torch.fx.traceback.preserve_node_meta():
+            interp = DebugInterpreter(fx_g)
+            out = interp.boxed_run(args)
+            mylogs = interp.get_logs()
+            if numerics_logger:
+                numerics_logger.log_fw_intermediates(mylogs)
+            return out
+
+    run._boxed_call = True
+    return run

examples/example_ds3_local_map.py

@@ -133,7 +133,12 @@ def input_fn():
             device=device,
         )
 
-    with AutoParallel(model, input_fn, mesh, dynamic=True) as autop:
+    numerics_logger = None
+    if rng_seed is not None:
+        numerics_logger = NumericsLogger(logs_dir)
+    with AutoParallel(
+        model, input_fn, mesh, dynamic=True, numerics_logger=numerics_logger
+    ) as autop:
         autop.add_parameter_memory_constraint(low=None, high=None)
 
         # x_sharding = (Shard(0), Replicate())
@@ -153,7 +158,7 @@ def input_fn():
     # )  # maybe not correct value
     parallel_mod.init_weights(buffer_device=device, seed=rng_seed)
     if rng_seed is not None:
-        NumericsLogger(logs_dir).log_model_weights(parallel_mod)
+        numerics_logger.log_model_weights(parallel_mod)
 
     x = (
         torch.randint(
@@ -173,7 +178,7 @@ def input_fn():
             shape_env=shape_env,
         ):
             # # now let's run it
-            out = parallel_mod(*x)
+            out = parallel_mod(*x, numerics_logger=numerics_logger)
             out.backward(torch.randn_like(out))
     else:
         out = parallel_mod(*x)