integrate torch.compile, controllable via config

run_llama_train.sh

@@ -9,4 +9,4 @@ NGPU=8
 MP=4
 
 torchrun --nproc_per_node=${NGPU} \
-train.py --steps 10
+train.py --steps 10 --compile

torchtrain/profiling.py

@@ -12,7 +12,7 @@
 
 from torchtrain.logging_utils import rank0_log
 
-_config_file = "./torchtrain/train_config.toml"
+_config_file = "./torchtrain/train_config/train_config.toml"
 
 
 def get_config_from_toml(config_path: str = _config_file) -> dict:

torchtrain/train_configs/__init_.py

@@ -0,0 +1 @@
+# __init__.py

train.py

@@ -2,8 +2,6 @@
 import os
 from dataclasses import dataclass, field
 from typing import List
-import logging
-from logging import getLogger
 
 # torch imports
 import torch
@@ -74,6 +72,14 @@ def main(args):
     optimizer = build_optimizer(model, args)
 
     # TODO: add metrics
+
+    # torch.compile model for improved performance
+    if args.compile:
+        rank0_log(f"Compiling model {model_name} with torch.compile...")
+        model = torch.compile(
+            model,
+        )
+
     train_state = TrainState()
 
     # train loop