From b5fd38715cda174a4263943d376180b5683f1a15 Mon Sep 17 00:00:00 2001 From: Masahiro Tanaka Date: Wed, 23 Oct 2024 05:13:32 +0000 Subject: [PATCH] remove pre/post process for profiling --- deepspeed/runtime/zero/compile/stage3_backend.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/deepspeed/runtime/zero/compile/stage3_backend.py b/deepspeed/runtime/zero/compile/stage3_backend.py index 52a9bec396a5..808784696ec5 100644 --- 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")