Utils Refactor (#180)

Co-authored-by: Philip Bontrager <[email protected]>

Author: pbontrager

recipes/__init__.py

@@ -5,8 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 
-_RECIPE_LIST = ["finetune_llm"]
-_CONFIG_LISTS = {"finetune_llm": ["alpaca_llama2_finetune"]}
+_RECIPE_LIST = ["finetune_llm", "llama_generate"]
+_CONFIG_LISTS = {"finetune_llm": ["alpaca_llama2_finetune"], "llama_generate": []}
 
 
 def list_recipes():

recipes/configs/alpaca_llama2_finetune.yaml

@@ -5,7 +5,7 @@ shuffle: True
 
 # Model Arguments
 model: llama2_7b
-model_checkpoint: /tmp/llama2-7b-01112024
+model_checkpoint: /tmp/llama2-7b
 tokenizer: llama2_tokenizer
 tokenizer_checkpoint: /tmp/tokenizer.model
 
@@ -17,5 +17,6 @@ optimizer: SGD
 loss: CrossEntropyLoss
 output_dir: /tmp/alpaca-llama2-finetune
 device: cuda
+dtype: fp32
 fsdp: False
 activation_checkpointing: False

recipes/finetune_llm.py

@@ -4,143 +4,111 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+
 import os
-import sys
 from functools import partial
-from typing import Callable
 
 import torch
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
-    apply_activation_checkpointing,
-)
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.distributed.fsdp.wrap import ModuleWrapPolicy
-from torch.optim.optimizer import Optimizer
+from torch.cuda.amp import GradScaler
 from torch.utils.data import DataLoader, DistributedSampler
 
-from torchtune.datasets import get_dataset, list_datasets
-from torchtune.models import get_model, get_tokenizer, list_models, list_tokenizers
-from torchtune.modules import TransformerDecoderLayer
-from torchtune.utils import TuneArgumentParser
-from torchtune.utils.batch_pad_sequence import batch_pad_to_longest_seq
-from torchtune.utils.env import get_world_size_and_rank, init_from_env, seed
+from torchtune import datasets, losses, models, modules, optim, utils
 from torchtune.utils.generation import generate_from_prompt
-from torchtune.utils.precision import (
-    get_autocast_manager,
-    get_grad_scaler,
-    get_supported_dtypes,
-)
 from tqdm import tqdm
 
 
-def get_optimizer(model: torch.nn.Module, optimizer: str, lr: float) -> Optimizer:
-    return getattr(torch.optim, optimizer)(model.parameters(), lr=lr)
-
-
-def get_loss(loss_fn: str) -> Callable:
-    return getattr(torch.nn, loss_fn)()
-
-
-def get_logger():
-    import logging
-
-    logger = logging.getLogger(__name__)
-    logger.addHandler(logging.StreamHandler())
-    logger.setLevel(logging.DEBUG)
-    return logger.info
-
-
-def recipe(kwargs):
+def recipe(
+    device,
+    dtype,
+    seed,
+    model,
+    model_checkpoint,
+    tokenizer,
+    tokenizer_checkpoint,
+    dataset,
+    shuffle,
+    batch_size,
+    fsdp,
+    epochs,
+    optimizer,
+    loss,
+    lr,
+    activation_checkpointing,
+    output_dir,
+    run_generation,
+    max_steps_per_epoch,
+):
     # ---- Initialize components ---- #
-    logger = get_logger()
-
-    # ---- Initialize distributed process group ---- #
-    device = init_from_env(device_type=kwargs["device"])
-    # TODO: only supporting devices specified as "cpu", "cuda", or "cuda:n" currently
-    device_type = (
-        kwargs["device"]
-        if kwargs["device"] in ("cpu", "cuda")
-        else kwargs["device"].split(":")[0]
-    )
-
-    # ---- Initialize seed ---- #
-    # Fetch world size and rank after distributed process group initialization
-    world_size, rank = get_world_size_and_rank()
-    if kwargs["seed"] is not None:
-        # Ensure that seed is different per rank (and its dataloader workers)
-        seed(kwargs["seed"] + rank)
-
-    tokenizer = get_tokenizer(kwargs["tokenizer"], path=kwargs["tokenizer_checkpoint"])
-    logger(msg=f"Loaded tokenizer from {kwargs['tokenizer_checkpoint']}")
-
-    autocast_precision = kwargs.get("autocast_precision", None)
-    autocast_mgr = get_autocast_manager(
-        device_type=device_type, precision=autocast_precision
-    )
-    grad_scaler = get_grad_scaler(autocast_precision, fsdp=kwargs["fsdp"])
-
-    model = get_model(
-        kwargs["model"],
-        device,
-    )
-
-    if kwargs["fsdp"] or kwargs["activation_checkpointing"]:
-        auto_wrap_policy = ModuleWrapPolicy(
-            {TransformerDecoderLayer}
-        )  # TODO: remove model specific components
-    if kwargs["fsdp"]:
-        model = FSDP(
-            model,
-            auto_wrap_policy=auto_wrap_policy,
-            device_id=device,
-            param_init_fn=lambda m: m.to_empty(device=device, recurse=False),
+    utils.init_distributed(fsdp)
+
+    # logger = logging.getLogger()
+    # logger.setLevel(logging.DEBUG) # test
+    logger = utils.get_logger("DEBUG")
+
+    device = utils.get_device(device)
+    dtype = utils.get_dtype(dtype)
+    seed = utils.set_seed(seed)
+
+    # ---- Setup model and load checkpoint ---- #
+    tokenizer = models.get_tokenizer(tokenizer, path=tokenizer_checkpoint)
+    logger.info(msg=f"Loaded tokenizer from {tokenizer_checkpoint}")
+
+    model = models.get_model(model, device=device)
+    if fsdp:
+        # TODO: initialize models for distributed on meta or cpu device to avoid OOMs
+        model = utils.get_fsdp(
+            model=model,
+            device=device,
+            dtype=dtype,
+            strategy="FULL_SHARD",
+            auto_wrap_policy={modules.TransformerDecoderLayer},
         )
-    if kwargs["activation_checkpointing"]:
-        apply_activation_checkpointing(
-            model,
-            check_fn=lambda mod: isinstance(
-                mod, TransformerDecoderLayer
-            ),  # TODO: remove model specific components
-            auto_wrap_policy=auto_wrap_policy,
+    if activation_checkpointing:
+        utils.set_activation_checkpointing(
+            model, auto_wrap_policy={modules.TransformerDecoderLayer}
         )
 
-    loaded_ckpt = torch.load(
-        kwargs["model_checkpoint"], map_location="cpu", weights_only=True
-    )
+    loaded_ckpt = torch.load(model_checkpoint, map_location="cpu", weights_only=True)
     model.load_state_dict(loaded_ckpt)
-    logger(msg=f"Loaded model from {kwargs['model_checkpoint']}")
+    logger.info(msg=f"Loaded model from {model_checkpoint}")
 
-    opt = get_optimizer(model, kwargs["optimizer"], kwargs["lr"])
+    # ---- Setup optimization functions ---- #
+    opt = optim.get_optimizer(optimizer, model, lr)
     # TODO add lr schedule option
-    loss_fn = get_loss(kwargs["loss"])
+    loss_fn = losses.get_loss(loss)
+
+    autocast = utils.get_autocast(dtype, device)
+    if dtype == torch.float16:
+        grad_scaler = utils.get_gradient_scaler(fsdp=fsdp)
+    else:
+        grad_scaler = GradScaler(enabled=False)
 
     # ---- Load dataset, set up sampler, and dataloader ---- #
-    dataset = get_dataset(kwargs["dataset"], split="train", tokenizer=tokenizer)
+    world_size, rank = utils.get_world_size_and_rank()
+    ds = datasets.get_dataset(dataset, split="train", tokenizer=tokenizer)
     sampler = DistributedSampler(
-        dataset,
+        ds,
         num_replicas=world_size,
         rank=rank,
-        shuffle=kwargs["shuffle"],
+        shuffle=shuffle,
         seed=0,
     )
     dataloader = DataLoader(
-        dataset=dataset,
-        batch_size=kwargs["batch_size"],
+        dataset=ds,
+        batch_size=batch_size,
         sampler=sampler,
         collate_fn=partial(
-            batch_pad_to_longest_seq,
-            input_padding_idx=tokenizer.pad_id,
-            label_padding_idx=loss_fn.ignore_index,  # TODO support loss without ignore_index
+            utils.padded_collate,
+            padding_idx=tokenizer.pad_id,
+            ignore_idx=loss_fn.ignore_index,  # TODO support loss without ignore_index
         ),
     )
-    logger(msg=f"Loaded dataset {kwargs['dataset']}")
+    logger.info(msg=f"Loaded dataset {dataset}")
 
     # ---- Train loop ---- #
-    for epoch in range(kwargs["epochs"]):
-        # Need to set the epoch for changing sample ordering in each epoch
-        sampler.set_epoch(epoch)
+    for epoch in range(epochs):
+        sampler.set_epoch(epoch)  # distributed sampler requires set_epoch
         for idx, batch in enumerate(pbar := tqdm(dataloader)):
-            max_steps_per_epoch = kwargs.get("max_steps_per_epoch", None)
             if max_steps_per_epoch is not None and idx == max_steps_per_epoch:
                 break
             opt.zero_grad()
@@ -149,10 +117,7 @@ def recipe(kwargs):
             input_ids = input_ids.to(device)
             labels = labels.to(device)
 
-            # Note: context manager for autocast is only applied in forward pass.
-            # see https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html#adding-torch-autocast
-            # for more details.
-            with autocast_mgr:
+            with autocast:
                 logits = model(input_ids)
                 # Shift so that tokens < n predict n
                 logits = logits[..., :-1, :].contiguous()
@@ -168,15 +133,11 @@ def recipe(kwargs):
                 f"{epoch+1}|{idx+1}|Loss: {loss.item()}"
             )  # TODO: add terminal logger
 
-            if grad_scaler:
-                grad_scaler.scale(loss).backward()
-                grad_scaler.step(opt)
-                grad_scaler.update()
-            else:
-                loss.backward()
-                opt.step()
+            grad_scaler.scale(loss).backward()
+            grad_scaler.step(opt)
+            grad_scaler.update()
 
-            run_generation = kwargs.get("run_generation", None)
+            # --- TODO TEMPORARY EVAL Code ---- #
             if run_generation and idx % run_generation == 0:
                 # Log a sample generation for the instruction.
                 # Just using a hardcoded prompt for now
@@ -189,16 +150,14 @@ def recipe(kwargs):
                 generation_str, decoded_tokens = generate_from_prompt(
                     prompt=prompt, tokenizer=tokenizer, decoder=model
                 )
-                if (
-                    not torch.distributed.is_initialized()
-                    or torch.distributed.get_rank() == 0
-                ):
-                    logger(f"Generation tokens: {decoded_tokens}")
-                    logger(f"Generation: {generation_str}")
-
-        # Save checkpoint at end of each epoch (to be changed later)
-        os.makedirs(kwargs["output_dir"], exist_ok=True)
-        output_loc = f"{kwargs['output_dir']}/model_{epoch}.ckpt"
+                if rank == 0:
+                    logger.info(f"Generation tokens: {decoded_tokens}")
+                    logger.info(f"Generation: {generation_str}")
+            # --- TODO TEMPORARY EVAL Code Ends ---- #
+
+        # ---- Save checkpoint at end of each epoch (to be changed later) ---- #
+        os.makedirs(output_dir, exist_ok=True)
+        output_loc = f"{output_dir}/model_{epoch}.ckpt"
         torch.save(
             {
                 "epoch": epoch,
@@ -208,19 +167,19 @@ def recipe(kwargs):
             },
             output_loc,
         )
-        logger(
+        logger.info(
             msg=f"Model checkpoint of size {os.path.getsize(output_loc) >> 20}MB saved to {output_loc}"
         )
 
 
 if __name__ == "__main__":
-    parser = TuneArgumentParser(description="Fine-tune an LLM.")
+    parser = utils.TuneArgumentParser(description="Fine-tune an LLM.")
 
     # Dataset and DataLoader arguments
     parser.add_argument(
         "--dataset",
         type=str,
-        choices=list_datasets(),
+        choices=datasets.list_datasets(),
         help="Dataset name.",
     )
     parser.add_argument(
@@ -238,7 +197,7 @@ def recipe(kwargs):
     parser.add_argument(
         "--model",
         type=str,
-        choices=list_models(),
+        choices=models.list_models(),
         help="Model to finetune.",
     )
     parser.add_argument(
@@ -249,7 +208,7 @@ def recipe(kwargs):
     parser.add_argument(
         "--tokenizer",
         type=str,
-        choices=list_tokenizers(),
+        choices=models.list_tokenizers(),
         help="Model tokenizer.",
     )
     parser.add_argument(
@@ -318,14 +277,12 @@ def recipe(kwargs):
         help="Max number of steps per epoch for faster dev/testing. Default is to finetune through the full dataset.",
     )
     parser.add_argument(
-        "--autocast-precision",
+        "--dtype",
         type=str,
-        choices=get_supported_dtypes(),
+        choices=utils.list_dtypes(),
         default=None,
-        help=f"""Low precision used for CUDA automatic mixed precision.
-            If specified, must be one of {get_supported_dtypes()}.
-        """,
+        help="Tensor dtype used for finetuning, lower precision types result in mixed precision training.",
     )
 
     kwargs = vars(parser.parse_args())
-    sys.exit(recipe(kwargs))
+    recipe(**kwargs)