Fix custom imports for both distributed and single device

torchtune/_cli/run.py

@@ -80,20 +80,29 @@ def _add_arguments(self) -> None:
             self._parser._add_action(action)
 
     @record
-    def _run_distributed(self, args: argparse.Namespace):
+    def _run_distributed(self, args: argparse.Namespace, is_builtin: bool):
         """Run a recipe with torchrun."""
         # TODO (rohan-varma): Add check that nproc_per_node <= cuda device count. Currently,
         # we don't do this since we test on CPUs for distributed. Will update once multi GPU CI is supported.
         print("Running with torchrun...")
         # Have to reset the argv so that the recipe can be run with the correct arguments
         args.training_script = args.recipe
         args.training_script_args = args.recipe_args
+        # torchtune built-in recipes are specified with an absolute posix path, but
+        # custom recipes are specified as a relative module dot path and need to be
+        # run with python -m
+        args.module = not is_builtin
         run(args)
 
-    def _run_single_device(self, args: argparse.Namespace):
+    def _run_single_device(self, args: argparse.Namespace, is_builtin: bool):
         """Run a recipe on a single device."""
         sys.argv = [str(args.recipe)] + args.recipe_args
-        runpy.run_path(str(args.recipe), run_name="__main__")
+        if is_builtin:
+            # torchtune built-in recipes are specified with an absolute posix path
+            runpy.run_path(str(args.recipe), run_name="__main__")
+        else:
+            # custom recipes are specified as a relative module dot path
+            runpy.run_module(str(args.recipe), run_name="__main__")
 
     def _is_distributed_args(self, args: argparse.Namespace):
         """Check if the user is trying to run a distributed recipe."""
@@ -155,9 +164,11 @@ def _run_cmd(self, args: argparse.Namespace):
         recipe = self._get_recipe(args.recipe)
         if recipe is None:
             recipe_path = args.recipe
+            is_builtin = False
         else:
             recipe_path = str(ROOT / "recipes" / recipe.file_path)
             supports_distributed = recipe.supports_distributed
+            is_builtin = True
 
         # Get config path
         config = self._get_config(config_str, recipe)
@@ -171,8 +182,6 @@ def _run_cmd(self, args: argparse.Namespace):
         args.recipe_args[config_idx] = config_path
 
         # Make sure user code in current directory is importable
-        # TODO: This is a temporary fix, figure out how to make runpy and torchrun
-        # run from this directory
         sys.path.append(os.getcwd())
 
         # Execute recipe
@@ -182,6 +191,6 @@ def _run_cmd(self, args: argparse.Namespace):
                     f"Recipe {recipe.name} does not support distributed training."
                     "Please run without torchrun commands."
                 )
-            self._run_distributed(args)
+            self._run_distributed(args, is_builtin=is_builtin)
         else:
-            self._run_single_device(args)
+            self._run_single_device(args, is_builtin=is_builtin)

torchtune/config/_instantiate.py

@@ -5,6 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 import copy
+import os
+import sys
 from typing import Any, Callable, Dict, Tuple
 
 from omegaconf import DictConfig, OmegaConf
@@ -89,6 +91,10 @@ def instantiate(
     if not OmegaConf.is_dict(config):
         raise ValueError(f"instantiate only supports DictConfigs, got {type(config)}")
 
+    # Ensure local imports are able to be instantiated
+    if os.getcwd() not in sys.path:
+        sys.path.append(os.getcwd())
+
     config_copy = copy.deepcopy(config)
     config_copy._set_flag(
         flags=["allow_objects", "struct", "readonly"], values=[True, False, False]