Fix circular import for MM dataprep notebook (#9287)

* update launcher name and fix mm circular import

Signed-off-by: Chen Cui <chcui@nvidia.com>

* Apply isort and black reformatting

Signed-off-by: cuichenx <cuichenx@users.noreply.github.com>

---------

Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: cuichenx <cuichenx@users.noreply.github.com>
Co-authored-by: cuichenx <cuichenx@users.noreply.github.com>

nemo/collections/nlp/parts/nlp_overrides.py

@@ -64,7 +64,6 @@
     # since PyTorch 2.3 the path has changed
     from torch.amp.grad_scaler import _refresh_per_optimizer_state
 
-from nemo.collections.multimodal.modules.stable_diffusion.attention import BasicTransformerBlock
 from nemo.collections.nlp.modules.common.megatron.module import Float16Module
 from nemo.collections.nlp.modules.common.megatron.transformer import AutocastTransformerLayer, ParallelTransformerLayer
 from nemo.collections.nlp.parts import utils_funcs
@@ -662,6 +661,9 @@ def __init__(
         # Use the default FSDP backward-prefetch policy for proper communication overlap.
         kwargs['backward_prefetch'] = BackwardPrefetch.BACKWARD_PRE
 
+        # import here to prevent circular imports
+        from nemo.collections.multimodal.modules.stable_diffusion.attention import BasicTransformerBlock
+
         # Set FSDP wrapping policy: use Transformer layer module as the FSDP sharding granularity.
         self.fsdp_wrap_module = {
             MCoreTransformerLayer,

tutorials/multimodal/Multimodal Data Preparation.ipynb

@@ -67,7 +67,7 @@
     "\n",
     "This notebook will show you how to prepare an image-text dataset into the [WebDataset](https://github.com/webdataset/webdataset) format. The Webdataset format is required to train all multimodal models in NeMo, such as Stable Diffusion and Imagen. \n",
     "\n",
-    "This notebook is designed to demonstrate the different stages of multimodal dataset preparation. It is not meant to be used to process large-scale datasets since many stages are too time-consuming to run without parallelism. For large workloads, we recommend running the multimodal dataset preparation pipeline with the NeMo-Megatron-Launcher on multiple processors/GPUs. NeMo-Megatron-Launcher packs the same 5 scripts in this notebook into one runnable command and one config file to enable a smooth and a streamlined workflow.\n",
+    "This notebook is designed to demonstrate the different stages of multimodal dataset preparation. It is not meant to be used to process large-scale datasets since many stages are too time-consuming to run without parallelism. For large workloads, we recommend running the multimodal dataset preparation pipeline with the NeMo-Framework-Launcher on multiple processors/GPUs. NeMo-Framework-Launcher packs the same 5 scripts in this notebook into one runnable command and one config file to enable a smooth and a streamlined workflow.\n",
     "\n",
     "Depending on your use case, not all 5 stages need to be run. Please go to [NeMo Multimodal Documentation](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/multimodal/text2img/datasets.html) for an overview of the 5 stages.\n",
     "    \n",
@@ -85,7 +85,7 @@
    "source": [
     "import os\n",
     "\n",
-    "LAUNCHER_DIR = \"/opt/NeMo-Megatron-Launcher\"\n",
+    "LAUNCHER_DIR = \"/opt/NeMo-Framework-Launcher\"  # formerly NeMo-Megatron-Launcher\n",
     "SCRIPT_DIR = os.path.join(LAUNCHER_DIR, \"launcher_scripts/nemo_launcher/collections/dataprep_scripts/multimodal_dataprep\")\n",
     "CONF_DIR = \"conf\"\n",
     "DATA_DIR = \"dummy_data\"\n",
@@ -168,7 +168,7 @@
     "\n",
     "Script: download_images.py\n",
     "\n",
-    "Environment variables (automatically set by SLURM if running with NeMo-Megatron-Launcher):\n",
+    "Environment variables (automatically set by SLURM if running with NeMo-Framework-Launcher):\n",
     "- `SLURM_ARRAY_TASK_COUNT`: total number of tasks, should be set to the number of parquet files in `$DATA_DIR/parquet/dummy_dataset50000.parquet_parts`. (i.e. `parquet_subpartitions` x `num_parquets_downloaded`)\n",
     "- `SLURM_ARRAY_TASK_ID`: id of the current task (0 <= SLURM_ARRAY_TASK_ID < SLURM_ARRAY_TASK_COUNT)\n",
     "\n",
@@ -266,7 +266,7 @@
     "\n",
     "Script: reorganize_tar.py\n",
     "\n",
-    "Environment variables (automatically set by SLURM if running with NeMo-Megatron-Launcher):\n",
+    "Environment variables (automatically set by SLURM if running with NeMo-Framework-Launcher):\n",
     "- `SLURM_ARRAY_TASK_COUNT`: total number of tasks, should be set to parquet_subpartitions x num_parquets_downloaded\n",
     "- `SLURM_ARRAY_TASK_ID`: id of the current task (0 <= `SLURM_ARRAY_TASK_ID` < `SLURM_ARRAY_TASK_COUNT`)\n",
     "\n",
@@ -430,7 +430,7 @@
    },
    "outputs": [],
    "source": [
-    "! wget https://raw.githubusercontent.com/NVIDIA/NeMo-Megatron-Launcher/master/launcher_scripts/conf/data_preparation/multimodal/precache_sd.yaml -P $CONF_DIR/"
+    "! wget https://raw.githubusercontent.com/NVIDIA/NeMo-Framework-Launcher/master/launcher_scripts/conf/data_preparation/multimodal/precache_sd.yaml -P $CONF_DIR/"
    ]
   },
   {
@@ -506,7 +506,7 @@
     "\n",
     "Script: precache_encodings.py\n",
     "\n",
-    "Environment variables (automatically set by SLURM if running with NeMo-Megatron-Launcher):\n",
+    "Environment variables (automatically set by SLURM if running with NeMo-Framework-Launcher):\n",
     "- `SLURM_ARRAY_TASK_COUNT`: total number of tasks, should be set to parquet_subpartitions x num_parquets_downloaded\n",
     "- `SLURM_ARRAY_TASK_ID`: id of the current task (0 <= `SLURM_ARRAY_TASK_ID` < `SLURM_ARRAY_TASK_COUNT`)\n",
     "\n",
@@ -533,15 +533,6 @@
     "    precache_config_path=$CONF_DIR/precache_sd_example.yaml"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "If you encounter a nemo import problem with the cell above, please also running it in the terminal directly."
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
   {
    "attachments": {},
    "cell_type": "markdown",