Merge pull request #22 from aws/release-1-2-1

julianhr · web-flow · commit f9f37ea9ecea · 2025-02-17T11:42:33.000-08:00
Sagemaker Hyperpod Recipes Release 1.2.1
diff --git a/launcher/nemo/k8s_templates/training/train-script-gpu.yaml b/launcher/nemo/k8s_templates/training/train-script-gpu.yaml
@@ -15,6 +15,7 @@ data:
     git clone {{ $config.git.repo_url_or_path }} $GIT_CLONE_DIR
     GIT_CLONE_DIR=${GIT_CLONE_DIR}/
     cd $GIT_CLONE_DIR
+    rm -rf __pycache__
 
       {{- if $config.git.branch }}
     git checkout {{ $config.git.branch }}
@@ -24,6 +25,11 @@ data:
     git fetch origin {{ $config.git.commit }}
     git reset --hard {{ $config.git.commit }}
       {{- end }}
+      {{- if $config.git.update_adapter }}
+
+    pip install . --force-reinstall --no-deps
+
+      {{- end }}
     {{- else }}
     GIT_CLONE_DIR=""
     {{- end }}
diff --git a/launcher/nemo/k8s_templates/training/train-script-trn.yaml b/launcher/nemo/k8s_templates/training/train-script-trn.yaml
@@ -17,6 +17,7 @@ data:
     git clone {{ $config.git.repo_url_or_path }} $GIT_CLONE_DIR
     GIT_CLONE_DIR=${GIT_CLONE_DIR}/
     cd $GIT_CLONE_DIR
+    rm -rf __pycache__
 
       {{- if $config.git.branch }}
     git checkout {{ $config.git.branch }}
@@ -26,6 +27,11 @@ data:
     git fetch origin {{ $config.git.commit }}
     git reset --hard {{ $config.git.commit }}
       {{- end }}
+      {{- if $config.git.update_adapter }}
+
+    pip install . --force-reinstall --no-deps
+
+      {{- end }}
     {{- else }}
     GIT_CLONE_DIR=""
     {{- end }}
diff --git a/launcher/nemo/k8s_templates/training/values.yaml b/launcher/nemo/k8s_templates/training/values.yaml
@@ -70,6 +70,7 @@ trainingConfig:
     branch: null
     commit: null
     token: null
+    update_adapter: null
 
   # Commands to run before training
   pre_script: []
diff --git a/launcher/nemo/stages.py b/launcher/nemo/stages.py
@@ -386,6 +386,9 @@ def _make_train_script_text(self, stage_cfg_path=None, port=41000) -> str:
                         f"git clone {repo_url_or_path} $GIT_CLONE_DIR",
                         "GIT_CLONE_DIR=${GIT_CLONE_DIR}/",
                         "cd $GIT_CLONE_DIR",
+                        # cache can lead to unexpected behavior when user clones
+                        # the Adapter and modifies it
+                        "rm -rf __pycache__",
                     ]
                 )
             else:
@@ -397,6 +400,8 @@ def _make_train_script_text(self, stage_cfg_path=None, port=41000) -> str:
             if self.cfg.get("git", None) is not None and self.cfg.git.get("commit", None) is not None:
                 script_text.append(f"git fetch origin {self.cfg.git.commit}")
                 script_text.append(f"git reset --hard {self.cfg.git.commit}")
+            if OmegaConf.select(self.cfg, "git.update_adapter", default=False):
+                script_text.append("\npip install . --force-reinstall --no-deps")
         else:
             script_text.append('GIT_CLONE_DIR=""')
 
@@ -703,6 +708,8 @@ def generate_default_k8s_value_template(self, template_root, cluster_parameters,
                 values_template.trainingConfig.git.branch = self.cfg.git.branch
             if self.cfg.git.get("commit", None) is not None:
                 values_template.trainingConfig.git.commit = self.cfg.git.commit
+            if self.cfg.git.get("update_adapter", None) is not None:
+                values_template.trainingConfig.git.update_adapter = self.cfg.git.update_adapter
 
         values_template.trainingConfig.device = self.device
         values_template.trainingConfig.scriptArgs = self.get_script_args_str(stage_cfg_path)
diff --git a/recipes_collection/config.yaml b/recipes_collection/config.yaml
@@ -27,6 +27,7 @@ git:
   commit: null
   entry_script: null
   token: null
+  update_adapter: false # if true it will re-install the Adapter code but not its dependencies
 
 env_vars:
   NCCL_DEBUG: WARN  # Logging level for NCCL. Set to "INFO" for debug information
diff --git a/tests/k8s_workflow/k8s_baseline_artifacts/llama-8b/k8s_template/values.yaml b/tests/k8s_workflow/k8s_baseline_artifacts/llama-8b/k8s_template/values.yaml
@@ -27,6 +27,7 @@ trainingConfig:
     branch: test_branch
     commit: test_commit
     token: null
+    update_adapter: false
   pre_script: []
   post_script: []
   labelSelector:
diff --git a/tests/k8s_workflow/k8s_baseline_artifacts/test_custom/k8s_template/values.yaml b/tests/k8s_workflow/k8s_baseline_artifacts/test_custom/k8s_template/values.yaml
@@ -26,6 +26,7 @@ trainingConfig:
     branch: null
     commit: null
     token: null
+    update_adapter: null
   pre_script: []
   post_script: []
   labelSelector:
diff --git a/tests/slurm_workflow/slurm_baseline_artifacts/hf-llama3-8b/train_script.sh b/tests/slurm_workflow/slurm_baseline_artifacts/hf-llama3-8b/train_script.sh
@@ -25,6 +25,7 @@ GIT_CLONE_DIR="$HOME/tmp/$LAUNCHER_HOSTNAME"
 git clone https://github.com/aws/sagemaker-hyperpod-training-adapter-for-nemo.git $GIT_CLONE_DIR
 GIT_CLONE_DIR=${GIT_CLONE_DIR}/
 cd $GIT_CLONE_DIR
+rm -rf __pycache__
 
 unset SLURM_NTASKS
 
diff --git a/tests/slurm_workflow/slurm_baseline_artifacts/llama-8b/train_script.sh b/tests/slurm_workflow/slurm_baseline_artifacts/llama-8b/train_script.sh
@@ -25,6 +25,7 @@ GIT_CLONE_DIR="$HOME/tmp/$LAUNCHER_HOSTNAME"
 git clone https://github.com/aws/sagemaker-hyperpod-training-adapter-for-nemo.git $GIT_CLONE_DIR
 GIT_CLONE_DIR=${GIT_CLONE_DIR}/
 cd $GIT_CLONE_DIR
+rm -rf __pycache__
 
 unset SLURM_NTASKS
 
diff --git a/tests/slurm_workflow/slurm_baseline_artifacts/test_custom/train_script.sh b/tests/slurm_workflow/slurm_baseline_artifacts/test_custom/train_script.sh
@@ -25,6 +25,7 @@ GIT_CLONE_DIR="$HOME/tmp/$LAUNCHER_HOSTNAME"
 git clone https://github.com/example $GIT_CLONE_DIR
 GIT_CLONE_DIR=${GIT_CLONE_DIR}/
 cd $GIT_CLONE_DIR
+rm -rf __pycache__
 
 unset SLURM_NTASKS
 
diff --git a/tests/test_readme.py b/tests/test_readme.py
@@ -0,0 +1,63 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+import logging
+from difflib import SequenceMatcher
+from typing import List
+
+from launcher.nemo.constants import ROOT_DIR
+
+logger = logging.getLogger(__name__)
+
+
+def test_readme_table_links():
+    readme_path = ROOT_DIR / "README.md"
+    log_line = lambda line: logger.info(f"\nFailing line:\n{line}")
+
+    def pluck_path_strings(line: str):
+        paths_str: List[str] = []
+
+        for chunk in line.split("|"):  # split by column delimeter
+            if "[link]" in chunk:
+                chunk = chunk.strip()
+                chunk = chunk.replace("[link]", "")
+                assert chunk[0] == "(" and chunk[-1] == ")", log_line(line)
+                chunk = chunk[1:-1]  # remove parantheses
+                paths_str.append(chunk)
+
+        return paths_str
+
+    with open(readme_path, "r") as fd:
+        for line in fd:
+            """
+            Example:
+            | Hugging Face | Llama 3.2 | 11b  | 8192            | 4     | ml.p5.48xlarge   | GPU H100    | [link](recipes_collection/recipes/training/llama/hf_llama3_2_11b_seq8k_gpu_p5x4_pretrain.yaml) | [link](launcher_scripts/llama/run_hf_llama3_2_11b_seq8k_gpu_p5x4_pretrain.sh) |
+            """
+            if "[link]" in line:
+                paths_str = pluck_path_strings(line)
+
+                if len(paths_str) == 1:
+                    file_path = ROOT_DIR / paths_str[0]
+                    assert file_path.exists(), log_line(line)
+                # there is a config and a script link
+                elif len(paths_str) == 2:
+                    config_file_path = ROOT_DIR / paths_str[0]
+                    launcher_script_path = ROOT_DIR / paths_str[1]
+                    # try to catch if a launch script is pointing to an incorrect config
+                    str_distance_ratio = SequenceMatcher(None, config_file_path.stem, launcher_script_path.stem).ratio()
+
+                    assert config_file_path.exists(), log_line(line)
+                    assert launcher_script_path.exists(), log_line(line)
+                    assert str_distance_ratio >= 0.8, log_line(line)
+                else:
+                    raise Exception("test condition not covered")
diff --git a/tests/test_recipes.py b/tests/test_recipes.py
@@ -1,8 +1,27 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
 import logging
 
+from omegaconf import OmegaConf
+
 from launcher.nemo.constants import ROOT_DIR
 
-from .test_utils import get_launcher_run_script_paths
+from .test_utils import (
+    get_launcher_run_script_paths,
+    make_hydra_cfg_instance,
+    validate_distributed_degrees,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -11,6 +30,9 @@
 
 def test_config_for_run_script_exists():
     RECIPES_DIR = ROOT_DIR / "recipes_collection/recipes"
+    log_line = lambda script, config: logger.info(
+        f"\nlauncher file: {script.relative_to(ROOT_DIR)}" f"\nconfig file: {config.relative_to(ROOT_DIR)}" "\n"
+    )
 
     for run_script_path in RUN_SCRIPT_PATHS:
         with open(run_script_path, "r") as fd:
@@ -25,10 +47,28 @@ def test_config_for_run_script_exists():
                     _, config_path_str = line.split("=")
                     config_path = RECIPES_DIR / (config_path_str + ".yaml")  # append .yaml
 
-                    logger.info(
-                        f"\nlauncher file: {run_script_path.relative_to(ROOT_DIR)}"
-                        f"\nconfig file: {config_path.relative_to(ROOT_DIR)}"
-                        "\n"
-                    )
+                    assert config_path.exists(), log_line(run_script_path, config_path)
+
+
+def test_config_degree_validation():
+    recipes_dir = ROOT_DIR / "recipes_collection/recipes"
+    log_config_name = lambda name: logger.info(f"\nFailing Config File: {name}")
+
+    for path in recipes_dir.rglob("*.yaml"):
+        if not path.is_file():
+            continue
+
+        # Hydra requires relative path definition
+        file_path: str = "../" + str(path.relative_to(ROOT_DIR).parent)
+        config = make_hydra_cfg_instance(file_path, path.name)
+
+        # plucking values outside the method arguments substantially reduces log output on failure
+        shard_degree = OmegaConf.select(config, "model.shard_degree")
+        tensor_model_parallel_degree = OmegaConf.select(config, "model.tensor_model_parallel_degree")
+        expert_model_parallel_degree = OmegaConf.select(config, "model.expert_model_parallel_degree")
+        context_parallel_degree = OmegaConf.select(config, "model.context_parallel_degree")
+        num_nodes = OmegaConf.select(config, "trainer.num_nodes")
 
-                    assert config_path.exists()
+        assert validate_distributed_degrees(
+            shard_degree, tensor_model_parallel_degree, expert_model_parallel_degree, context_parallel_degree, num_nodes
+        ), log_config_name(path.name)
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -1,19 +1,36 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
 import difflib
 import logging
 import os
 import shutil
 import stat
 import tempfile
+from functools import lru_cache
 from pathlib import Path
-from typing import List
+from typing import Any, List, Optional
 
 from hydra import compose, initialize
 from hydra.core.hydra_config import HydraConfig
+from omegaconf import DictConfig
 
 from launcher.nemo.constants import ROOT_DIR
 
 logger = logging.getLogger(__name__)
 
+GPUS_PER_NODE = 8
+
 
 def create_temp_directory():
     """Create a temporary directory and Set full permissions for the directory"""
@@ -72,21 +89,22 @@ def compare_files(file1_path, file2_path):
     return True
 
 
-def compose_hydra_cfg(path, config_name, overrides=[]):
+def compose_hydra_cfg(config_path: str, config_name: str, overrides: List[Any] = []) -> DictConfig:
     """Init and compose a hydra config"""
-    with initialize(version_base=None, config_path=path):
+    with initialize(version_base=None, config_path=config_path):
         return compose(config_name=config_name, overrides=overrides, return_hydra_config=True)
 
 
-def make_hydra_cfg_instance(path, config_name, overrides):
+def make_hydra_cfg_instance(config_path: str, config_name: str, overrides=[]) -> DictConfig:
     """Init hydra instance"""
     # Note: This is needed if using compose API and not hydra.main b/c we rely on hydra resolver
     # Open issue tracking fix https://github.com/facebookresearch/hydra/issues/2017
-    config = compose_hydra_cfg(path, config_name, overrides)
+    config = compose_hydra_cfg(config_path, config_name, overrides)
     HydraConfig.instance().set_config(config)
     return config
 
 
+@lru_cache(maxsize=None)
 def get_launcher_run_script_paths() -> List[Path]:
     scripts_dir = ROOT_DIR / "launcher_scripts"
     launch_script_paths = []
@@ -99,3 +117,33 @@ def get_launcher_run_script_paths() -> List[Path]:
             launch_script_paths.append(path)
 
     return launch_script_paths
+
+
+# adapted from
+# https://github.com/aws/sagemaker-hyperpod-training-adapter-for-nemo/blob/83c6037ad566ccb2584e7becab01da2ca8b6af20/src/hyperpod_nemo_adapter/conf/config_schemas.py#L78
+def validate_distributed_degrees(
+    shard_degree: Optional[int],
+    tensor_model_parallel_degree: Optional[int],
+    expert_model_parallel_degree: Optional[int],
+    context_parallel_degree: Optional[int],
+    num_nodes: Optional[int],
+) -> bool:
+    """
+    Check that the degrees are legal.
+    """
+    # default param values to 1 if they are missing
+    sd = shard_degree or 1
+    tp = tensor_model_parallel_degree or 1
+    ep = expert_model_parallel_degree or 1
+    cp = context_parallel_degree or 1
+    degree_mult = sd * tp * ep
+    gpu_per_node = int(os.environ.get("LOCAL_WORLD_SIZE", GPUS_PER_NODE))
+    world_size = (num_nodes or 1) * gpu_per_node
+
+    # Validate the degree multiplication <= world_size
+    degree_mult_is_valid = world_size % degree_mult == 0
+
+    # Validate CP degree <= shard degree
+    cp_is_valid = cp <= sd if cp > 1 else True
+
+    return degree_mult_is_valid and cp_is_valid
