integrate torch.compile, controllable via config

torchtrain/profiling.py

@@ -5,28 +5,14 @@
 import os
 import torch
 
-try:
-    import tomllib
-except ModuleNotFoundError:
-    import tomli as tomllib
 
 from torchtrain.logging_utils import rank0_log
-
-_config_file = "./torchtrain/train_config.toml"
-
-
-def get_config_from_toml(config_path: str = _config_file) -> dict:
-    """
-    Reads a config file in TOML format and returns a dictionary.
-    """
-    with open(config_path, "rb") as f:
-        config = tomllib.load(f)
-    return config
+from torchtrain.tt_config.config_utils import get_config
 
 
 @contextlib.contextmanager
 def maybe_run_profiler(*pos_args, **kwargs):
-    config = get_config_from_toml()
+    config = get_config()
 
     # get user defined profiler settings
     run_profiler = config["profiling"].get("run_profiler", False)

torchtrain/tt_config/__init_.py

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

torchtrain/tt_config/config_utils.py

@@ -0,0 +1,13 @@
+try:
+    import tomllib
+except ModuleNotFoundError:
+    import tomli as tomllib
+
+
+def get_config(config_path: str = "./torchtrain/tt_config/train_config.toml") -> dict:
+    """
+    Reads a config file in TOML format and returns a dictionary.
+    """
+    with open(config_path, "rb") as f:
+        config = tomllib.load(f)
+    return config

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

@@ -2,6 +2,9 @@
 [global]
 dump_folder = "./torchtrain/outputs"
 
+[compile]
+use_compile = true
+
 [profiling]
 run_profiler = true
 save_traces_folder = "profiling/traces"

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
@@ -24,6 +21,8 @@
     pad_batch_to_longest_seq,
 )
 
+from torchtrain.tt_config.config_utils import get_config
+
 
 @dataclass
 class TrainState:
@@ -56,6 +55,10 @@ def main(args):
         device_type, (dp_degree, args.tp_degree), mesh_dim_names=("dp", "tp")
     )
 
+    # load config
+    tt_config = get_config()
+    _use_compile = tt_config["compile"]["use_compile"]
+
     model_name = args.model
     # build tokenizer
     tokenizer_type = model_name_to_tokenizer[model_name]
@@ -85,6 +88,14 @@ def main(args):
 
     # TODO: apply parallelisms, e.g. fsdp/tp
     # TODO: add metrics
+
+    # torch.compile model for improved performance
+    if _use_compile:
+        rank0_log(f"Compiling model {model_name} with torch.compile...")
+        torch.compile(
+            model,
+        )
+
     train_state = TrainState()
 
     # train loop