remove pre/post process for profiling

tohtana · Oct 23, 2024 · b5fd387 · b5fd387
1 parent 9cf9f88
commit b5fd387
Showing 1 changed file with 0 additions and 4 deletions.
diff --git a/deepspeed/runtime/zero/compile/stage3_backend.py b/deepspeed/runtime/zero/compile/stage3_backend.py
@@ -237,10 +237,8 @@ def fw(gm, sample_inputs):
             gm.recompile()
 
             if debug_log:
-                nz3.start_forward()
                 mem_prof = MemoryProfilingInterpreter(gm)
                 mem_prof.run(*real_inputs)
-                nz3.end_forward()
                 if rank == 0:
                     mem_prof.dump(f"mem_prof_fwd_{graph_id}.csv")
 
@@ -310,10 +308,8 @@ def bw(gm, sample_inputs):
             gm.recompile()
 
             if debug_log:
-                nz3.start_backward(True)
                 mem_prof = MemoryProfilingInterpreter(gm)
                 mem_prof.run(*validated_inputs)
-                nz3.end_backward()
                 if rank == 0:
                     mem_prof.dump(f"mem_prof_bwd_{graph_id}.csv")