allow float types

configs/mcli/small-test.yaml

@@ -18,14 +18,10 @@ command: |-
 
   pip freeze
 
-  # Prepare environment including AWS config files for both S3 and R2 access.
+  # Prepare environment.
   mkdir -p /root/.cache/torch
+  export OMP_NUM_THREADS=8
+  export LOG_FILTER_TYPE=local_rank0_only
 
-  torchrun \
-  --master_addr "$MASTER_ADDR" \
-  --master_port "$MASTER_PORT" \
-  --nnodes "$NUM_NODES" \
-  --node_rank "$NODE_RANK" \
-  --nproc_per_node 8 \
-  scripts/train.py configs/small-test-s3.yaml \
+  torchrun --nproc_per_node 8 scripts/train.py configs/small-test-s3.yaml \
     --run_name=${run_name}

olmo/config.py

@@ -489,11 +489,11 @@ class SchedulerUnits(StrEnum):
 class SchedulerConfig(BaseConfig):
     name: SchedulerType = SchedulerType.cosine_with_warmup
     units: SchedulerUnits = SchedulerUnits.steps
-    t_warmup: int = 100
-    t_max: Optional[int] = None
+    t_warmup: Union[int, float] = 100
+    t_max: Optional[Union[int, float]] = None
     alpha_f: float = 0.1
 
-    grad_clip_warmup_steps: Optional[int] = None
+    grad_clip_warmup_steps: Optional[Union[int, float]] = None
     """
     The warmup period for which the max grad norm (or norm ratio) will be set to its
     warmup value of `max_grad_norm * grad_clip_warmup_factor`.