Enable memory profiling for D2Go trainer of Genie

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/607

Titled

Reviewed By: tglik

Differential Revision: D47535500

fbshipit-source-id: 93635f36b7164472bac6560d9f6626262096d14e

d2go/runner/config_defaults.py

@@ -18,7 +18,10 @@ from d2go.quantization.modeling import add_quantization_default_configs
 from d2go.registry.builtin import CONFIG_UPDATER_REGISTRY
 from d2go.trainer.activation_checkpointing import add_activation_checkpoint_configs
 from d2go.trainer.fsdp import add_fsdp_configs
-from d2go.utils.gpu_memory_profiler import add_memory_profiler_configs
+from d2go.utils.gpu_memory_profiler import (
+    add_memory_profiler_configs,
+    add_zoomer_default_config,
+)
 from d2go.utils.visualization import add_tensorboard_default_configs
 from detectron2.config import get_cfg as get_d2_cfg
 from mobile_cv.common.misc.oss_utils import fb_overwritable
@@ -138,6 +141,9 @@ def _add_detectron2go_runner_default_cfg(_C: CN) -> None:
     # GPU memory profiler
     add_memory_profiler_configs(_C)
 
+    # Zoomer memory profiling
+    add_zoomer_default_config(_C)
+
     # Checkpointing-specific config
     _C.LOAD_CKPT_TO_GPU = False
 

d2go/utils/gpu_memory_profiler.py

@@ -23,6 +23,14 @@ def add_memory_profiler_configs(_C: CN):
     _C.MEMORY_PROFILER.LOG_DURING_TRAIN_AT = 550
 
 
+def add_zoomer_default_config(_C: CN):
+    _C.ZOOMER = CN()
+    _C.ZOOMER.ENABLE_STACK_TRACING = (
+        False  # Do not enable by default, since it may cause performance regression
+    )
+    _C.ZOOMER.ENABLE_MEMORY_PROFILING = False
+
+
 def omm_logger_wrapper(output_dir):
     def oom_logger(
         device: int, alloc: int, device_alloc: int, device_free: int