Trainer + reference model in bf16 (#189)

Author: joecummings

apps/grpo/qwen3_1_7b.yaml

@@ -46,7 +46,8 @@ trainer:
     local_batch_size: ${batch_size}
     seq_len: 2048
     max_norm: 1.0
-    steps: 5
+    steps: 1000000
+    dtype: bfloat16
   compile:
     enable: false
   parallelism:
@@ -80,6 +81,8 @@ ref_model:
     name: qwen3
     flavor: 1.7B
     hf_assets_path: hf://${model}
+  training:
+    dtype: bfloat16
   compile:
     enable: false
   parallelism:

src/forge/actors/reference_model.py

@@ -15,7 +15,13 @@
 from monarch.actor import current_rank, current_size, endpoint
 from torch.distributed.tensor import DTensor
 
-from torchtitan.config.job_config import Checkpoint, Compile, Model, Parallelism
+from torchtitan.config.job_config import (
+    Checkpoint,
+    Compile,
+    Model,
+    Parallelism,
+    Training,
+)
 from torchtitan.experiments.forge.engine import ForgeEngine
 from torchtitan.experiments.forge.job_config import ForgeJobConfig
 
@@ -29,6 +35,9 @@ class ReferenceModel(ForgeActor):
     parallelism: Parallelism = field(default_factory=Parallelism)
     checkpoint: Checkpoint = field(default_factory=Checkpoint)
     compile: Compile = field(default_factory=Compile)
+    training: Training = field(
+        default_factory=Training
+    )  # Only needed in order to correctly set a lower dtype
 
     # Populated in setup
     # TODO: Commented out since engine_config parsing extracts from class members