adjust hidden size, NCCL log level

allenai · Dec 3, 2024 · 355f61a · 355f61a
1 parent 3dfdedf
commit 355f61a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/src/olmo_core/internal/common.py b/src/olmo_core/internal/common.py
@@ -6,6 +6,7 @@
 from olmo_core.io import is_url
 from olmo_core.launch.beaker import (
     BeakerEnvSecret,
+    BeakerEnvVar,
     BeakerLaunchConfig,
     BeakerWekaBucket,
     OLMoCoreBeakerImage,
@@ -61,6 +62,7 @@ def build_launch_config(
     task_name: str = "train",
     workspace: str = "ai2/OLMo-core",
     budget: str = "ai2/oe-training",
+    nccl_debug: bool = False,
 ) -> BeakerLaunchConfig:
     weka_buckets: List[BeakerWekaBucket] = []
     if root_dir.startswith("/weka/"):
@@ -81,6 +83,7 @@ def build_launch_config(
         num_gpus=8,
         shared_filesystem=not is_url(root_dir),
         allow_dirty=False,
+        env_vars=[BeakerEnvVar(name="NCCL_DEBUG", value="INFO" if nccl_debug else "WARN")],
         env_secrets=[
             BeakerEnvSecret(name="BEAKER_TOKEN", secret=f"{beaker_user}_BEAKER_TOKEN"),
             BeakerEnvSecret(name="WANDB_API_KEY", secret=f"{beaker_user}_WANDB_API_KEY"),

diff --git a/src/olmo_core/nn/transformer/config.py b/src/olmo_core/nn/transformer/config.py
@@ -460,6 +460,7 @@ def olmo2_26B(cls, vocab_size: int, **kwargs) -> "TransformerConfig":
             block_name=kwargs.pop("block_name", TransformerBlockType.reordered_norm),
             qk_norm=kwargs.pop("qk_norm", True),
             rope_theta=kwargs.pop("rope_theta", 500_000),
+            hidden_size_multiple_of=kwargs.pop("hidden_size_multiple_of", 1024),
             layer_norm_eps=1e-6,
             **kwargs,
         )