Optim package (#457)

stanford-crfm · Feb 7, 2024 · e034003 · e034003
1 parent 378470a
commit e034003
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/docs/Configuration-Guide.md b/docs/Configuration-Guide.md
@@ -13,7 +13,7 @@ class TrainLmConfig:
     data: LMDatasetConfig = field(default_factory=LMDatasetConfig)
     trainer: TrainerConfig = field(default_factory=TrainerConfig)
     model: LmConfig = field(default_factory=Gpt2Config)
-    optimizer: OptimizerConfig = field(default_factory=OptimizerConfig)
+    optimizer: OptimizerConfig = field(default_factory=AdamConfig)
 ```
 
 Your training run will typically be associated with a single config file. For instance, you might have a file
@@ -290,7 +290,7 @@ If you're not using SLURM or TPUs, you can specify the cluster manually using th
 
 ## Optimizer
 
-[levanter.trainer.OptimizerConfig][] is a dataclass that specifies the optimizer configuration. It has the following fields:
+[levanter.optim.OptimizerConfig][] is a dataclass that specifies the optimizer configuration. It has the following fields:
 
 | Parameter       | Description                                                       | Default  |
 |-----------------|-------------------------------------------------------------------|----------|
@@ -358,7 +358,7 @@ trainer:
 
 ### Optimizer
 
-::: levanter.trainer.OptimizerConfig
+::: levanter.optim.OptimizerConfig
 
 ### LM Model
 

diff --git a/src/levanter/trainer.py b/src/levanter/trainer.py
@@ -5,7 +5,6 @@
 import os
 import sys
 import typing
-import warnings
 from dataclasses import dataclass
 from functools import cached_property
 from pathlib import Path
@@ -51,7 +50,7 @@
 from levanter.distributed import DistributedConfig, RayConfig
 from levanter.grad_accum import microbatched
 from levanter.logging import capture_time
-from levanter.optim import SecondOrderTransformation
+from levanter.optim import OptimizerConfig, SecondOrderTransformation  # noqa: F401
 from levanter.tracker import TrackerConfig
 from levanter.types import FilterSpec
 from levanter.utils import cloud_utils