From 146c4eb02335819cd80399d9812f0350d4dde0f1 Mon Sep 17 00:00:00 2001 From: ranzhejiang Date: Mon, 18 Nov 2024 07:05:09 +0000 Subject: [PATCH 1/2] Makes the with_stack of the profiler changeable --- optimum/habana/transformers/trainer.py | 1 + optimum/habana/transformers/training_args.py | 5 +++++ optimum/habana/utils.py | 3 ++- 3 files changed, 8 insertions(+), 1 deletion(-) diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 27eab623bb..07e2f077a7 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -881,6 +881,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio warmup=self.args.profiling_warmup_steps, active=self.args.profiling_steps, record_shapes=self.args.profiling_record_shapes, + with_stack=self.args.profiling_with_stack, ) hb_profiler.start() diff --git a/optimum/habana/transformers/training_args.py b/optimum/habana/transformers/training_args.py index 44af85fd54..2a5ddc8b4d 100644 --- a/optimum/habana/transformers/training_args.py +++ b/optimum/habana/transformers/training_args.py @@ -242,6 +242,11 @@ class GaudiTrainingArguments(TrainingArguments): default=True, metadata={"help": ("Record shapes when enabling profiling.")}, ) + + profiling_with_stack: Optional[bool] = field( + default=False, + metadata={"help": ("record source information (file and line number) for the ops when enabling profiling.")}, + ) # Overriding the default value of optim because 'adamw_hf' is deprecated optim: Optional[Union[OptimizerNames, str]] = field( default="adamw_torch", diff --git a/optimum/habana/utils.py b/optimum/habana/utils.py index a16c36b3fb..a3aa728e81 100755 --- a/optimum/habana/utils.py +++ b/optimum/habana/utils.py @@ -285,6 +285,7 @@ def __init__( warmup: int = 0, active: int = 0, record_shapes: bool = True, + with_stack: bool = False, output_dir: str = "./hpu_profile", wait: int = 0, ): @@ -306,7 +307,7 @@ def noop(): activities=activities, on_trace_ready=torch.profiler.tensorboard_trace_handler(output_dir), record_shapes=record_shapes, - with_stack=False, + with_stack=with_stack, ) self.start = profiler.start self.stop = profiler.stop From 9d066e2716d25b0dbbacc19f0a59b7d614779c65 Mon Sep 17 00:00:00 2001 From: ranzhejiang Date: Fri, 22 Nov 2024 03:51:59 +0000 Subject: [PATCH 2/2] format the blank line --- optimum/habana/transformers/training_args.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/habana/transformers/training_args.py b/optimum/habana/transformers/training_args.py index 2a5ddc8b4d..286601158d 100644 --- a/optimum/habana/transformers/training_args.py +++ b/optimum/habana/transformers/training_args.py @@ -242,7 +242,7 @@ class GaudiTrainingArguments(TrainingArguments): default=True, metadata={"help": ("Record shapes when enabling profiling.")}, ) - + profiling_with_stack: Optional[bool] = field( default=False, metadata={"help": ("record source information (file and line number) for the ops when enabling profiling.")},