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_config/__init_.py

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

torchtrain/train_config.toml → torchtrain/train_config/train_config.toml

@@ -3,7 +3,7 @@
 dump_folder = "./torchtrain/outputs"
 
 [profiling]
-run_profiler = true
+run_profiler = false
 save_traces_folder = "profiling/traces"
 # profiling frequency - example: 10 means every 10th iter will be profiled
 profile_every_x_iter = 10

train.py

@@ -2,9 +2,6 @@
 import os
 from dataclasses import dataclass, field
 from typing import List
-import logging
-from logging import getLogger
-import sys  # for logging
 
 # torch imports
 import torch
@@ -56,6 +53,7 @@ def main(args):
         device_type, (dp_degree, args.tp_degree), mesh_dim_names=("dp", "tp")
     )
 
+
     model_name = args.model
     # build tokenizer
     tokenizer_type = model_name_to_tokenizer[model_name]
@@ -85,6 +83,15 @@ def main(args):
 
     # TODO: apply parallelisms, e.g. fsdp/tp
     # 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