Add memory profiler arguments to Python frontend

Author: tbennun

applications/nlp/transformer/trainer.py

@@ -172,9 +172,7 @@ def make_batch_script(model: lbann.Model,
     if args.progress:
         model.callbacks.append(lbann.CallbackProgressBar(newline_interval=100))
 
-    profiler = lbann.contrib.args.create_profile_callback(args)
-    if profiler is not None:
-        model.callbacks.append(profiler)
+    model.callbacks.extend(lbann.contrib.args.create_profile_callbacks(args))
 
     script_params = lbann.contrib.args.get_scheduler_kwargs(args)
     script_params['work_dir'] = work_dir

applications/physics/cosmology/cosmoflow/train_cosmoflow.py

@@ -201,10 +201,8 @@ def create_synthetic_data_reader(input_width: int, num_responses: int) -> Any:
                                                       mlperf=args.mlperf,
                                                       transform_input=args.transform_input)
 
-    # Add profiling callback if needed.
-    profiler = lbann.contrib.args.create_profile_callback(args)
-    if profiler is not None:
-        model.callbacks.append(profiler)
+    # Add profiling callbacks if needed.
+    model.callbacks.extend(lbann.contrib.args.create_profile_callbacks(args))
 
     # Setup optimizer
     optimizer = lbann.contrib.args.create_optimizer(args)

applications/vision/resnet.py

@@ -136,9 +136,7 @@
     callbacks.append(
         lbann.CallbackLinearGrowthLearningRate(
             target=0.1 * args.mini_batch_size / 256, num_epochs=5))
-profiler = lbann.contrib.args.create_profile_callback(args)
-if profiler is not None:
-    callbacks.append(profiler)
+callbacks.extend(lbann.contrib.args.create_profile_callbacks(args))
 model = lbann.Model(args.num_epochs,
                     layers=layers,
                     objective_function=obj,

python/lbann/contrib/args.py

@@ -209,8 +209,8 @@ def add_profiling_arguments(parser: argparse.ArgumentParser) -> None:
     """Add command-line arguments for common profiler settings within
     LBANN.
 
-    Adds the following options: `--profile`, `--profile-init`,
-    `--caliper` and `--caliper-config`.
+    Adds the following options: `--profile`, `--memory-profile`,
+    `--profile-init`, `--caliper` and `--caliper-config`.
 
     `--caliper-config` implies `--caliper`. `--caliper` without a
     `--caliper-config` will use the default configuration in LBANN.
@@ -229,6 +229,10 @@ def add_profiling_arguments(parser: argparse.ArgumentParser) -> None:
                         action='store_true',
                         default=False,
                         help='enable profiling instrumentation and markers')
+    parser.add_argument('--memory-profile',
+                        action='store_true',
+                        default=False,
+                        help='enable itemized memory usage analysis')
     parser.add_argument('--profile-init',
                         action='store_true',
                         default=False,
@@ -242,9 +246,12 @@ def add_profiling_arguments(parser: argparse.ArgumentParser) -> None:
                         default=None,
                         type=str,
                         help='Configuration string for Caliper')
+    parser.add_argument('--memory-profile-verbose',
+                        action='store_true',
+                        default=False,
+                        help='increase memory usage analysis verbosity')
 
-
-def create_profile_callback(args: argparse.Namespace) -> Any:
+def create_profile_callbacks(args: argparse.Namespace) -> Any:
     """Create a profiler callback from command-line arguments.
 
     The parsed arguments must be generated by an
@@ -262,13 +269,20 @@ def create_profile_callback(args: argparse.Namespace) -> Any:
     try:
         profile = args.profile
         profile_init = not args.profile_init
+        memprofile = args.memory_profile
+        memprof_verbose = args.memory_profile_verbose
     except AttributeError:
         raise ValueError('passed arguments have not been processed by '
                          '`add_profiling_arguments`')
 
+    result = []
     if profile:
-        return lbann.CallbackProfiler(skip_init=profile_init)
-    return None
+        result.append(lbann.CallbackProfiler(skip_init=profile_init))
+    if memprofile:
+        result.append(lbann.CallbackMemoryProfiler(
+            detailed_first_step=memprof_verbose))
+
+    return result
 
 
 def get_profile_args(args: argparse.Namespace) -> list[str]: