Merge pull request #65 from aws/release-2.0.0

Sagemaker Hyperpod Recipes Release 2.0.0

Author: aws-tianquaw

README.md

@@ -0,0 +1,139 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+from hydra.core.config_store import ConfigStore
+
+### LLMFT hydra schema
+
+
+@dataclass
+class LLMFTTrainerConfig:
+    devices: int = 1
+    num_nodes: int = 1
+
+
+@dataclass
+class LLMFTTrainingArgsConfig:
+    micro_train_batch_size: int = 1
+    train_batch_size: int = 16
+    learning_rate: float = 0.0001
+    lr_warmup_ratio: float = 0.1
+    gradient_clipping: bool = True
+    gradient_clipping_threshold: float = 1.0
+    max_epochs: int = 3
+    logging_steps: int = 1
+    save_steps: int = 0
+    eval_steps: int = -1
+    beta: float = 0.01
+    nll_loss_coef: float = 0.0
+    label_smoothing: float = 0.0
+
+
+@dataclass
+class LLMFTrainingConfig:
+    training_args: LLMFTTrainingArgsConfig = field(default_factory=LLMFTTrainingArgsConfig)
+
+
+@dataclass
+class LLMFTRecipeConfig:
+    trainer: LLMFTTrainerConfig = field(default_factory=LLMFTTrainerConfig)
+    training_config: Optional[LLMFTrainingConfig] = field(default=None)
+
+
+### NOVA hydra schema
+@dataclass
+class NovaModelConfig:
+    hidden_dropout: float = 0.0
+    attention_dropout: float = 0.0
+    ffn_dropout: float = 0.0
+
+
+@dataclass
+class NovaTrainerConfig:
+    max_epochs: int = 1
+
+
+@dataclass
+class NovaTrainingConfig:
+    max_length: int = 8192
+    global_batch_size: int = 256
+    trainer: NovaTrainerConfig = field(default_factory=NovaTrainerConfig)
+    model: NovaModelConfig = field(default_factory=NovaModelConfig)
+
+
+@dataclass
+class NovaRecipeConfig:
+    training_config: NovaTrainingConfig = field(default_factory=NovaTrainingConfig)
+
+
+### Verl hydra schema
+@dataclass
+class VerlModelOptimConfig:
+    """Model optimizer configuration for VERL."""
+
+    lr: float = 1e-5
+
+
+@dataclass
+class VerlModelConfig:
+    """Model configuration for VERL."""
+
+    path: Optional[str] = None
+
+
+@dataclass
+class VerlCriticConfig:
+    """Critic configuration for VERL."""
+
+    optim: VerlModelOptimConfig = field(default_factory=VerlModelOptimConfig)
+    model: VerlModelConfig = field(default_factory=VerlModelConfig)
+    ppo_micro_batch_size_per_gpu: int = 4
+
+
+@dataclass
+class VerlKlCtrlConfig:
+    """KL control configuration for VERL."""
+
+    kl_coef: float = 0.001
+    target_kl: float = 0.1
+
+
+@dataclass
+class VerlAlgorithmConfig:
+    """Algorithm configuration for VERL."""
+
+    kl_ctrl: VerlKlCtrlConfig = field(default_factory=VerlKlCtrlConfig)
+    adv_estimator: Optional[str] = None  # "gae" or "grpo"
+
+
+@dataclass
+class VerlRayInitConfig:
+    """Ray initialization configuration for VERL."""
+
+    num_cpus: Optional[int] = None
+    timeline_json_file: Optional[str] = None
+
+
+@dataclass
+class VerlRecipeConfig:
+    """Top-level configuration for VERL recipes."""
+
+    critic: Optional[VerlCriticConfig] = field(default_factory=VerlCriticConfig)
+    algorithm: Optional[VerlAlgorithmConfig] = field(default_factory=VerlAlgorithmConfig)
+    ray_init: Optional[VerlRayInitConfig] = field(default_factory=VerlRayInitConfig)
+
+
+# Register with Hydra
+cs = ConfigStore.instance()
+cs.store(name="recipe_schema", node=LLMFTRecipeConfig())
+cs.store(group="trainer", name="base_trainer", node=LLMFTTrainerConfig())
+cs.store(group="training_config", name="base_training", node=LLMFTrainingConfig())
+
+# Register Nova configs with Hydra
+cs.store(name="nova_recipe_schema", node=NovaRecipeConfig())
+cs.store(group="training_config", name="nova_training", node=NovaTrainingConfig())
+
+# Register Verl configs with Hydra
+cs.store(name="verl_recipe_schema", node=VerlRecipeConfig())
+cs.store(group="algorithm", name="verl_algorithm", node=VerlAlgorithmConfig())
+cs.store(group="critic", name="verl_critic", node=VerlCriticConfig())

docs/RECIPES.md

@@ -0,0 +1,54 @@
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+
+### LLMFT Validators
+class LLMFTTrainerValidator(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    devices: int = Field(gt=0)
+    num_nodes: int = Field(gt=0)
+
+
+class LLMFTTrainingArgsValidator(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    micro_train_batch_size: int | None = Field(default=None, gt=0)
+    train_batch_size: int | None = Field(default=None, gt=0)
+    learning_rate: float | None = Field(default=None, gt=0)
+    lr_warmup_ratio: float | None = Field(default=None, ge=0, le=1)
+    gradient_clipping: bool | None = None
+    gradient_clipping_threshold: float | None = Field(default=None, gt=0)
+    max_epochs: int | None = Field(default=None, gt=0)
+    logging_steps: int | None = Field(default=None, gt=0)
+    save_steps: int | None = Field(default=None, ge=0)
+    eval_steps: int | None = Field(default=None, ge=-1)
+    beta: float | None = Field(default=None, gt=0)
+    nll_loss_coef: float | None = Field(default=None, ge=0)
+    label_smoothing: float | None = Field(default=None, ge=0, le=1)
+
+    @model_validator(mode="after")
+    def check_batch_sizes(self):
+        if self.micro_train_batch_size and self.train_batch_size:
+            if self.train_batch_size < self.micro_train_batch_size:
+                raise ValueError("train_batch_size must be >= micro_train_batch_size")
+        return self
+
+
+class LLMFTRecipeValidator(BaseModel):
+    """Top-level validator for LLMFT recipes."""
+
+    model_config = ConfigDict(extra="allow")
+
+    trainer: LLMFTTrainerValidator | None = None
+    training_config: dict | None = None
+
+    @model_validator(mode="after")
+    def validate_nested_fields(self):
+        # Validate training_config.training_args if present
+        if self.training_config and "training_args" in self.training_config:
+            try:
+                LLMFTTrainingArgsValidator(**self.training_config["training_args"])
+            except Exception as e:
+                raise ValueError(f"Error validating training_args: {str(e)}")
+
+        return self

launcher/config_validator/schema/base.py

@@ -0,0 +1,68 @@
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+
+### Verl Validators
+class VerlModelOptimValidator(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    lr: float | None = Field(gt=0)
+
+
+class VerlModelConfigValidator(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    path: str | None = None
+
+
+class VerlCriticValidator(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    optim: VerlModelOptimValidator | dict | None = None
+    model: VerlModelConfigValidator | dict | None = None
+    ppo_micro_batch_size_per_gpu: int | None = Field(gt=0)
+
+
+class VerlKlCtrlValidator(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    kl_coef: float | None = Field(default=0.001, gt=0)
+    target_kl: float | None = Field(default=0.1, gt=0)
+
+
+class VerlAlgorithmValidator(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    kl_ctrl: VerlKlCtrlValidator | None = None
+    adv_estimator: str | None = None  # gae or grpo
+
+    @model_validator(mode="after")
+    def validate_adv_estimator(self):
+        if self.adv_estimator and self.adv_estimator not in ["gae", "grpo"]:
+            raise ValueError("adv_estimator must be either 'gae' or 'grpo'")
+        return self
+
+
+class VerlRecipeValidator(BaseModel):
+    """Top-level validator for VERL recipes."""
+
+    model_config = ConfigDict(extra="allow")
+
+    critic: VerlCriticValidator | dict | None = None
+    algorithm: VerlAlgorithmValidator | dict | None = None
+    ray_init: dict | None = None
+
+    @model_validator(mode="after")
+    def validate_nested_fields(self):
+        # Validate algorithm if present
+        if self.algorithm and isinstance(self.algorithm, dict):
+            try:
+                VerlAlgorithmValidator(**self.algorithm)
+            except Exception as e:
+                raise ValueError(f"Error validating algorithm configuration: {str(e)}")
+        if self.critic and isinstance(self.critic, dict):
+            try:
+                VerlCriticValidator(**self.critic)
+            except Exception as e:
+                raise ValueError(f"Error validating critic configuration: {str(e)}")
+
+        return self

launcher/config_validator/schema/llmft_schema_validation.py

@@ -170,6 +170,7 @@
     ]
 )
 
+# Mapping of instance types to their GPU/device counts (8 is the default)
 INSTANCE_TO_DEVICE_COUNT = {
     "g4dn.xlarge": 1,
     "g4dn.2xlarge": 1,

launcher/config_validator/schema/verl_schema_validation.py

@@ -0,0 +1,12 @@
+# 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.

launcher/efa.py

@@ -0,0 +1,24 @@
+# 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.
+
+# Evaluation container constants
+EVAL_CONTAINER_IMAGE = "{account_id}.dkr.ecr.{region}.amazonaws.com/smdistributed-modelparallel:2.4.1-gpu-py311-cu121"
+
+# Region to account mapping for evaluation containers
+EVAL_REGION_ACCOUNT_MAP = {
+    "us-east-1": "658645717510",
+    "us-west-2": "658645717510",
+    "eu-west-1": "658645717510",
+    "ap-southeast-1": "658645717510",
+    "ap-northeast-1": "658645717510",
+}

launcher/evaluation/__init__.py

@@ -0,0 +1,6 @@
+apiVersion: v2
+name: evaluation-job
+description: A Helm chart for SageMaker HyperPod evaluation jobs
+type: application
+version: 0.1.0
+appVersion: "1.0"

launcher/evaluation/constants.py

@@ -0,0 +1,15 @@
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: evaluation-config-{{ .Values.evaluationConfig.jobName }}
+  {{- if .Values.evaluationConfig.namespace }}
+  namespace: {{ .Values.evaluationConfig.namespace }}
+  {{- end }}
+  labels:
+    app: {{ .Values.evaluationConfig.jobName }}
+    {{- if .Values.evaluationConfig.customLabels }}
+    {{- toYaml .Values.evaluationConfig.customLabels | nindent 4 }}
+    {{- end }}
+data:
+  recipe.yaml: |
+{{ (.Files.Get "config/{{ .Values.evaluationConfig.jobName }}_hydra.yaml") | indent 4 }}