Repository: instructlab/instructlab

Issue:
Training CI cleanup catch-all
Most items from review discussion on #1557 

- [ ] fix `FULL_TRAIN` so that training does not run twice on when FULL_TRAIN is true
- [ ] establish if re-installation of deepspeed is necessary and if it is a fix in the CLI or in the training repo
- [ ] remove packaging installation in the e2e workflow. Figure out why it is causing failures since it is a requirement in training repo
- [ ] update training-lib based training to do training based on the output of `ilab data generate` instead of a file checked in to the tree - `scripts/test-data/train_all_pruned_SDG.jsonl`
- [ ] update `docs/ci.md` to reflect where the training-lib is used
write full_train.py for CPU and MPS support
full_train.py should implement a custom training loop that by default trains the granite-7b-lab model using a memory efficient optimizer and the CPU/MPS torch device depending on what is available

This will work better for CPU since MPS is known to not be optimized on recent released of torch
remove `--legacy` flag and overhaul `linux_train.py`
once `full_train.py` is in place and MLX is redefined. the only linux train code should be removed since it should not be used anymore


Discussion:
@cdoern did https://github.com/instructlab/instructlab/pull/1557 fix this?
@nathan-weinberg no this should be handled as part of the CI cleanup that I think @alimaredia is working on







Patch:
diff --git a/.spellcheck-en-custom.txt b/.spellcheck-en-custom.txt
index 4e40197534..00e495076e 100644
--- a/.spellcheck-en-custom.txt
+++ b/.spellcheck-en-custom.txt
@@ -93,9 +93,11 @@ Nvidia
 NVidia
 OCI
 oneMKL
+OOM
 orchestrator
 ots
 parallelized
+PEFT
 Pareja
 PEFT
 PlantUML
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f63633106e..1d8916f29b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,7 @@
   * `test`
 * Intel Gaudi software has been updated to 1.17.1 with Python 3.11 and
   Torch 2.3.1 support.
+* `--legacy` has been removed and replaced with `--pipeline=simple` in `ilab model train`
 
 ### Features
 
@@ -24,6 +25,8 @@
 * `ilab model evaluate --max-workers=auto` is now supported and is the default option. When
   auto is specified, the optimal value is determined based on your GPUs, CPUs, and
   configuration.
+* `ilab model train` now supports `--pipeline`. The supported pipelines are `simple`, `full`, and `accelerated`. Simple preserves the functionality found in `--legacy` and current MacOS training. Full introduces a new training loop optimized for CPU and MacOS performance. `accelerated` allows users with dedicated graphics cards to run the full fine tuning and multi-phase training found in our training library.
+* `--device` in `ilab model train` now supports `mps` which stands for Metal Performance Shaders. This is a PyTorch device for MacOS training that allows us to utilize the same code path for Linux and MacOS.
 
 ## v0.18.1
 
diff --git a/docs/ci.md b/docs/ci.md
index 76d35a61d9..f8c576b3d8 100644
--- a/docs/ci.md
+++ b/docs/ci.md
@@ -46,20 +46,22 @@ the E2E job configuration files found
 | --- | --- |
 | `e` | Run model evaluation |
 | `F` | Run "fullsize" SDG |
-| `L` | Run legacy training |
+| `q` | Run the 'simple' training pipeline with 4 bit quantization |
+| `s` | Run the 'simple' training pipeline |
+| `f` | Use the 'full' training pipeline optimized for CPU and MPS rather than simple training |
+| `a` | Use the 'accelerated' training library rather than simple or full training |
 | `m` | Run minimal configuration (lower number of instructions and training epochs) |
 | `M` | Use Mixtral model (4-bit quantized) instead of Merlinite (4-bit quantized) |
 | `P` | Use the phased training within the 'full' training library |
-| `T` | Use the 'full' training library rather than legacy training |
 | `v` | Run with vLLM for serving |
 
 ### Current E2E Jobs
 
 | Name | T-Shirt Size | Runner Host | Instance Type | OS | GPU Type | Script flags | Runs when? | Slack reporting? |
 | --- | --- | --- | --- | --- | --- | --- | --- | --- |
-| [`e2e-nvidia-t4-x1.yml`](https://github.com/instructlab/instructlab/blob/main/.github/workflows/e2e-nvidia-t4-x1.yml) | Small | AWS | [`g4dn.2xlarge`](https://aws.amazon.com/ec2/instance-types/g4/) | CentOS Stream 9 | 1 x NVIDIA Tesla T4 w/ 16 GB VRAM | `m` | Pull Requests, Push to `main` or `release-*` branch | No |
-| [`e2e-nvidia-a10g-x1.yml`](https://github.com/instructlab/instructlab/blob/main/.github/workflows/e2e-nvidia-a10g-x1.yml) | Medium | AWS |[`g5.2xlarge`](https://aws.amazon.com/ec2/instance-types/g5/) | CentOS Stream 9 | 1 x NVIDIA A10G w/ 24 GB VRAM | `mf` | Manually by Maintainers | No |
-| [`e2e-nvidia-a10g-x4.yml`](https://github.com/instructlab/instructlab/blob/main/.github/workflows/e2e-nvidia-a10g-x4.yml) | Large | AWS |[`g5.12xlarge`](https://aws.amazon.com/ec2/instance-types/g5/) | CentOS Stream 9 | 4 x NVIDIA A10G w/ 24 GB VRAM (98 GB) | `mevFMT` | Manually by Maintainers, Automatically against `main` branch at 8AM/8PM UTC | Yes |
+| [`e2e-nvidia-t4-x1.yml`](https://github.com/instructlab/instructlab/blob/main/.github/workflows/e2e-nvidia-t4-x1.yml) | Small | AWS | [`g4dn.2xlarge`](https://aws.amazon.com/ec2/instance-types/g4/) | CentOS Stream 9 | 1 x NVIDIA Tesla T4 w/ 16 GB VRAM | `msq` | Pull Requests, Push to `main` or `release-*` branch | No |
+| [`e2e-nvidia-a10g-x1.yml`](https://github.com/instructlab/instructlab/blob/main/.github/workflows/e2e-nvidia-a10g-x1.yml) | Medium | AWS |[`g5.4xlarge`](https://aws.amazon.com/ec2/instance-types/g5/) | CentOS Stream 9 | 1 x NVIDIA A10G w/ 24 GB VRAM | `mf` | Manually by Maintainers | No |
+| [`e2e-nvidia-a10g-x4.yml`](https://github.com/instructlab/instructlab/blob/main/.github/workflows/e2e-nvidia-a10g-x4.yml) | Large | AWS |[`g5.12xlarge`](https://aws.amazon.com/ec2/instance-types/g5/) | CentOS Stream 9 | 4 x NVIDIA A10G w/ 24 GB VRAM (98 GB) | `mevFMa` | Manually by Maintainers, Automatically against `main` branch at 8AM/8PM UTC | Yes |
 
 ### E2E Test Coverage Matrix
 
@@ -69,15 +71,16 @@ the E2E job configuration files found
 |              | vllm                     |⎯|⎯|✅|
 | **Generate** | simple                   |✅|✅|⎯|
 |              | full                     |⎯|⎯|✅|
-| **Training** | legacy+Linux             |⎯|❌|⎯|
-|              | legacy+Linux+4-bit-quant |✅|⎯|⎯|
-|              | training-lib             |⎯|⎯|✅|
-| **Eval**     | eval                     |⎯|❌(*2)|✅|
+| **Training** | simple                   |✅(*1)|⎯|⎯|
+|              | full                     |⎯ |✅|⎯|
+|              | accelerated              |⎯|⎯|✅(*2)|
+| **Eval**     | eval                     |⎯|❌(*3)|✅|
 
 Points of clarification (*):
 
-1. The `training-lib` testing is not testing using the output of the Generate step. <https://github.com/instructlab/instructlab/issues/1655>
-2. The `eval` testing is not evaluating the output of the Training step. <https://github.com/instructlab/instructlab/issues/1642>
+1. The `simple` training pipeline uses 4-bit-quantization.
+2. The `accelerated` testing is not using the output of the Generate step. <https://github.com/instructlab/instructlab/issues/1655>
+3. The `eval` testing is not evaluating the output of the Training step. <https://github.com/instructlab/instructlab/issues/1642>
 
 ### Current E2E Status against `main`
 
diff --git a/requirements.txt b/requirements.txt
index 5576b64f35..afde39cd33 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -35,3 +35,4 @@ transformers>=4.41.2
 trl>=0.9.4
 wandb>=0.16.4
 xdg-base-dirs>=6.0.1
+psutil>=6.0.0
diff --git a/src/instructlab/configuration.py b/src/instructlab/configuration.py
index c209bdf4d1..3e03ce050c 100644
--- a/src/instructlab/configuration.py
+++ b/src/instructlab/configuration.py
@@ -1275,6 +1275,8 @@ def map_train_to_library(ctx, params):
 
     train_args.deepspeed_options = ds_args
     train_args.lora = lora_args
+    if params["pipeline"] == "full":
+        train_args.disable_flash_attn = True
 
     return train_args, torch_args
 
diff --git a/src/instructlab/model/full_train.py b/src/instructlab/model/full_train.py
new file mode 100644
index 0000000000..2c920b0860
--- /dev/null
+++ b/src/instructlab/model/full_train.py
@@ -0,0 +1,324 @@
+# Standard
+from copy import copy
+from functools import partial
+from pathlib import Path
+import logging
+import math
+import os
+
+# Third Party
+from instructlab_quantize import run_quantize
+from tqdm import tqdm
+from transformers import Adafactor, AutoModelForCausalLM
+import numpy as np
+import psutil
+
+# First Party
+from instructlab.llamacpp import llamacpp_convert_to_gguf
+
+logger = logging.getLogger(__name__)
+
+
+def train(train_args, device):
+    """
+    train runs a CPU and MacOS optimized version of full fine tuning.
+    Adafactor is the optimizer of choice and the multiprocessing method is set to spawn.
+    Dataloading functions imported from the training library.
+    """
+
+    # pylint: disable=no-name-in-module
+    # Third Party
+    from instructlab.training import config
+    from instructlab.training import data_process as dp
+    from instructlab.training import (
+        multipack_sampler,
+        token_dataset,
+        tokenizer_utils,
+        utils,
+    )
+    from torch.utils.data import DataLoader
+    import torch
+
+    dp.main(
+        config.DataProcessArgs(
+            data_output_path=train_args.data_output_dir,
+            model_path=train_args.model_path,
+            data_path=train_args.data_path,
+            max_seq_len=train_args.max_seq_len,
+            chat_tmpl_path=train_args.chat_tmpl_path,
+        )
+    )
+
+    # load chat template based on path in the args
+    CHAT_TEMPLATE, SPECIAL_TOKENS = utils.retrieve_chat_template(
+        train_args.chat_tmpl_path
+    )
+    # get the tokenizer for the model
+    tokenizer = tokenizer_utils.setup_tokenizer(
+        train_args.model_path, SPECIAL_TOKENS, CHAT_TEMPLATE
+    )
+
+    # setup the dataset and place it in data.jsonl, this needs to be used for training NOT the jsonl produced by sdg
+    dataset = token_dataset.setup_dataset(
+        os.path.join(train_args.data_output_dir, "data.jsonl"),
+    )
+
+    # based on the length of the dataset, figure out the max batch len
+    packing_max_batch_len, _ = (
+        multipack_sampler.find_packing_max_batch_len_and_grad_accum(
+            num_gpus=1,
+            avg_sample_len=dataset.get_lengths().mean(),
+            effective_batch_size=train_args.effective_batch_size,
+            max_batch_len_per_gpu=train_args.max_batch_len,
+            is_padding=train_args.is_padding_free,
+            pad_id=tokenizer.pad_token_id,
+            dataset=dataset,
+            seed=47,
+        )
+    )
+
+    collate_fn = partial(pad_collate_fn, pad_token_id=tokenizer.pad_token_id)
+
+    # use a multipack sampler since this is a non-dist. scenario, training library defaults to this as well
+    sampler = multipack_sampler.MultipackDistributedBatchSampler(
+        batch_max_length=packing_max_batch_len,
+        lengths=dataset.get_lengths(),
+        num_replicas=1,
+        rank=0,
+        seed=47,
+        padding=True,
+    )
+    sampler = {"batch_sampler": sampler}
+
+    # 4 workers for the dataloader as compared to the original 8 to optimize performance
+    # this dataloader needs a spawn multiproc method to optimize memory and to work on Apple Silicon.
+    dataloader = DataLoader(
+        dataset,
+        **sampler,
+        num_workers=4,
+        collate_fn=collate_fn,
+    )
+    dataloader.multiprocessing_context = "spawn"
+
+    # set device based off argument given
+    dev = torch.device(device)
+    # auto model based on model path
+    model = AutoModelForCausalLM.from_pretrained(train_args.model_path)
+    if len(tokenizer) > model.config.vocab_size:
+        logger.warning(
+            f"WARNING: tokenizer has {len(tokenizer)} tokens but model has {model.config.vocab_size} vocab size"
+        )
+        model.resize_token_embeddings(
+            int(8 * math.ceil(len(tokenizer) / 8.0))
+        )  # make the vocab size multiple of 8 for sharding the embedding layer.
+
+    # Fix any discrepancy between model and tokenizer
+    # bos and eos tokens mark the beginning and end of a sequence/sentence
+    if (
+        model.config.pad_token_id is not None
+        and tokenizer.pad_token_id is not None
+        and model.config.pad_token_id != tokenizer.pad_token_id
+    ):
+        logger.warning(
+            f"There is a mismatch between pad token id of model ({model.config.pad_token_id}) and tokenizer ({tokenizer.pad_token_id}). Fixing model pad token id to be same as tokenizer's pad token id."
+        )
+        model.config.pad_token_id = tokenizer.pad_token_id
+    if (
+        model.config.bos_token_id is not None
+        and tokenizer.bos_token_id is not None
+        and model.config.bos_token_id != tokenizer.bos_token_id
+    ):
+        logger.warning(
+            f"There is a mismatch between bos token id of model ({model.config.bos_token_id}) and tokenizer ({tokenizer.bos_token_id}). These tokens denote the start of a sequence of data. Fixing model bos token id to be same as tokenizer's bos token id."
+        )
+        model.config.bos_token_id = tokenizer.bos_token_id
+    if (
+        model.config.eos_token_id is not None
+        and tokenizer.eos_token_id
+        and model.config.eos_token_id != tokenizer.eos_token_id
+    ):
+        logger.warning(
+            f"There is a mismatch between eos token id of model ({model.config.eos_token_id}) and tokenizer ({tokenizer.eos_token_id}). These tokens denote the end of a sequence of data. Fixing model eos token id to be same as tokenizer's eos token id."
+        )
+        model.config.eos_token_id = tokenizer.eos_token_id
+
+    model = utils.convert_loss_to_reduce_sum(model, is_granite=False)
+    model = utils.add_noisy_embeddings(model, noise_alpha=None)
+
+    # Get virtual memory statistics
+    memory_info = psutil.virtual_memory()
+
+    # Total RAM
+    total_ram = memory_info.total / (1024**3)  # Convert to GB
+    logger.info(f"Total RAM: {total_ram:.2f} GB")
+    # if RAM is <= 16, we need to use fp16 not fp32. This will yield a worse model but will allow the full pipeline to run
+    accum = 1
+    if total_ram <= 16:
+        # if <= 16GB ram, use gradinent accum and hald precision
+        logger.warning(
+            f"Your system has {total_ram:.2f} GB of RAM. This is below our reccomendation of 32GB for this type of training. Using half precision."
+        )
+        model = model.to(dev).half()  # Convert model to float16
+        accum = 4
+    else:
+        model = model.to(dev)
+
+    # adafactor and gradient checkpointing are memory friendly, we opt to use these in the CPU/MPS loop to fit 7b models.
+    optimizer = Adafactor(
+        model.parameters(), lr=1e-5, scale_parameter=True, relative_step=False
+    )
+    model.gradient_checkpointing_enable()
+
+    model.train()
+
+    # For each epoch do the following:
+    # 1. Incremement Loading Bar
+    # 2. put zeros onto the proper device so we can store some values we need
+    # 3. put each part of the batch onto the device as well using fp6
+    # 4. Forward pass on the model. If using half, we need to do some special handling. If not, then just put the batch onto the device using fp32
+    # 5. Calculate the loss
+    # 6. backward pass on the loss
+    # 7. We are using gradient accumulation to be memory efficient, so if this step is divisible by 4, take a step and zero out the gradients (freeing memory)
+    # 8. clear out some caches if on MPS
+    # 9. Incrememnt loading bar
+    # 10. at the end og the epoch, save the checpoint to a samples folder and convert to GGUF as well
+    for epoch in range(train_args.num_epochs):
+        dataloader.batch_sampler.set_epoch(epoch)
+        inner_pb = tqdm(range(len(dataloader)), desc=f"Epoch {epoch}")
+        aggregated_values = torch.zeros(3, dtype=torch.float16).to(dev)
+
+        for step, batch in enumerate(dataloader):
+            aggregated_values[0] = batch.pop("num_loss_counted_tokens")
+            aggregated_values[1] = len(batch["input_ids"])
+
+            # Move and cast batch data to device
+            for k in batch:
+                if total_ram < 16:
+                    if k in ["input_ids", "attention_mask"]:
+                        # these two need to be of type long if using .half()
+                        batch[k] = batch[k].to(device=dev, dtype=torch.long)
+                    else:
+                        batch[k] = batch[k].to(device=dev, dtype=torch.float16)
+                else:
+                    batch[k] = batch[k].to(device=dev)
+
+            output = model(**batch, use_cache=False)
+            loss = output.loss
+            aggregated_values[2] = loss.item()
+
+            num_loss_counted_tokens = aggregated_values[0]
+            loss = loss / num_loss_counted_tokens
+
+            loss = loss / accum  # Scale the loss for accumulation steps
+
+            logger.info(f"\nEpoch: {epoch}, Step: {step + 1}, Rank: 0, loss = {loss}")
+
+            # Gradient accumulation
+            loss.backward()  # Backward pass
+
+            # Clear cache before optimizer step
+            # below we clear the MPS cache quite a bit. MPS caching is different from CPU.
+            # on MPS systems with 32 or 64 GB of unified memory, it is important to clear the cache beacuse often we are paging a lot of the data that can't fit directly on GPU
+            # CPU doesn't have an equivalent cache clearing method, only GPU devices like MPS, CUDA, ROCM, etc.
+            if dev.type == "mps":
+                torch.mps.empty_cache()
+
+            # if we are on a step which is divisible by 4, step and zero gradients
+            if (step + 1) % accum == 0:
+                optimizer.step()  # Optimizer step
+                optimizer.zero_grad()  # Zero gradients
+
+                # Clear cache after optimizer step
+                if dev.type == "mps":
+                    torch.mps.empty_cache()
+
+            inner_pb.update(1)
+
+        # Clear cache at the end of the epoch if needed
+        if dev.type == "mps":
+            torch.mps.empty_cache()
+
+        output_dir = (
+            Path(train_args.ckpt_output_dir) / "hf_format" / f"samples_{(epoch*8)}"
+        )
+
+        logger.info(f"Saving Model to: {output_dir}")
+        model_state = model.state_dict()
+        output_dir.mkdir(parents=True, exist_ok=True)
+        output_model_file = output_dir / "pytorch_model.bin"
+        output_config_file = output_dir / "config.json"
+
+        tmp_conf = copy(model.config)
+
+        torch.save(model_state, str(output_model_file))
+        tmp_conf.to_json_file(str(output_config_file))
+        tokenizer.save_pretrained(str(output_dir))
+
+        # convert to GGUF at the end so llamacpp can serve the model originally in .bin
+        llamacpp_convert_to_gguf.convert_llama_to_gguf(
+            model=output_dir,
+            pad_vocab=True,
+            skip_unknown=True,
+            outfile=os.path.join(output_dir, "pytorch_model.gguf"),
+        )
+
+        # quantize the model so everyone can run it
+        gguf_model_dir = os.path.join(output_dir, "pytorch_model.gguf")
+        gguf_model_q_dir = os.path.join(output_dir, "pytorch_model-Q4_K_M.gguf")
+        run_quantize(gguf_model_dir, gguf_model_q_dir, "Q4_K_M")
+
+
+def pad_collate_fn(batch, pad_token_id):
+    lens = np.array([len(item["input_ids"]) for item in batch])
+    max_len = max(lens)
+    # Third Party
+    import torch
+    import torch.nn.functional as F
+
+    input_ids = torch.stack(
+        [
+            F.pad(
+                item["input_ids"],
+                (max_len - len(item["input_ids"]), 0),
+                mode="constant",
+                value=pad_token_id,
+            )
+            for item in batch
+        ]
+    )
+    labels = torch.stack(
+        [
+            F.pad(
+                item["labels"],
+                (max_len - len(item["labels"]), 0),
+                mode="constant",
+                value=-100,
+            )
+            for item in batch
+        ]
+    )
+    num_loss_counted_tokens = (labels != -100).sum()
+
+    attention_mask = torch.stack(
+        [
+            F.pad(
+                item["attention_mask"],
+                (max_len - len(item["attention_mask"]), 0),
+                mode="constant",
+                value=0,
+            )
+            for item in batch
+        ]
+    )
+    logger.info(
+        f"\033[96m total tokens: {max_len * len(batch)} num samples: {len(batch)} num padding tokens: {max_len * len(batch) - lens.sum()} - rank: {0} "
+        f"max len: {max_len} min len: {min(lens)} avg len: {lens.mean()} "
+        f"num_loss_counted_tokens: {num_loss_counted_tokens}\033[0m"
+    )
+
+    return {
+        "input_ids": input_ids,
+        "labels": labels,
+        "num_loss_counted_tokens": num_loss_counted_tokens,
+        "attention_mask": attention_mask,
+    }
diff --git a/src/instructlab/model/train.py b/src/instructlab/model/train.py
index 0c9dc8f199..223044a493 100644
--- a/src/instructlab/model/train.py
+++ b/src/instructlab/model/train.py
@@ -124,12 +124,13 @@ def clickpath_setup(is_dir: bool) -> click.Path:
 )
 @click.option(
     "--device",
-    type=click.Choice(["cpu", "cuda", "hpu"]),
+    type=click.Choice(["cpu", "cuda", "hpu", "mps"]),
     show_default=True,
     default="cpu",
     help=(
-        "PyTorch device for Linux training. Use 'cuda' "
+        "PyTorch device for full training. Use 'cuda' "
         "for NVidia CUDA / AMD ROCm GPU, to use specific GPU, set visible GPU before run train command."
+        "use 'cpu' for laptop CPU training and 'mps' for MacOS Metal Performance Shader training."
     ),
 )
 @click.option(
@@ -276,11 +277,6 @@ def clickpath_setup(is_dir: bool) -> click.Path:
     config_sections=ADDITIONAL_ARGUMENTS,
     required=True,  # default from config
 )
-@click.option(
-    "--legacy",
-    is_flag=True,
-    help="if true, enables the legacy linux training code path from release 0.17.0 and prior.",
-)
 @click.option(
     "--strategy",
     type=click.Choice(
@@ -356,6 +352,14 @@ def clickpath_setup(is_dir: bool) -> click.Path:
     is_flag=True,
     help="By default, checkpoints are saved at the end of each training epoch. This option disables this behavior.",
 )
+@click.option(
+    "--pipeline",
+    type=click.Choice(["simple", "full", "accelerated"]),
+    default="accelerated",
+    help="Model fidelity pipeline for training: 'simple' uses SFTTrainer on Linux or MLX on MacOS, producing low-fidelity models quickly for rapid prototyping."
+    "'full' employs CPU and MPS optimized InstructLab fine-tuning, generating medium-fidelity models over a longer period."
+    "'accelerated' utilizes GPU acceleration and distributed training, yielding high-fidelity models but requiring more time. Choose based on your hardware, time constraints, and desired model quality",
+)
 @click.pass_context
 @clickext.display_params
 def train(
@@ -372,7 +376,6 @@ def train(
     num_epochs,
     device: str,
     four_bit_quant: bool,
-    legacy,
     strategy: str | None,
     phased_base_dir: pathlib.Path,
     phased_phase1_data: pathlib.Path | None,
@@ -386,12 +389,19 @@ def train(
     phased_mt_bench_judge: pathlib.Path | None,
     skip_user_confirm: bool,
     enable_serving_output: bool,
+    pipeline: str,
     **kwargs,
 ):
     """
     Takes synthetic data generated locally with `ilab data generate` and the previous model and learns a new model using the MLX API.
     On success, writes newly learned model to {model_dir}/mlx_model, which is where `chatmlx` will look for a model.
     """
+    if (
+        pipeline in ("full", "simple")
+        and strategy == SupportedTrainingStrategies.LAB_MULTIPHASE.value
+    ):
+        ctx.fail("Multi Phase training is only supported with `--pipeline accelerated`")
+
     if not input_dir:
         # By default, generate output-dir is used as train input-dir
         input_dir = ctx.obj.config.generate.output_dir
@@ -463,218 +473,238 @@ def get_files(directory: str, pattern: str) -> list[str]:
         shutil.copy(train_files[0], train_file)
         shutil.copy(test_files[0], test_file)
 
-    # if macos, preserve that path
-    if utils.is_macos_with_m_chip() and not strategy:
-        # Local
-        from ..mlx_explore.gguf_convert_to_mlx import load
-        from ..mlx_explore.utils import fetch_tokenizer_from_hub
-        from ..train.lora_mlx.convert import convert_between_mlx_and_pytorch
-        from ..train.lora_mlx.lora import load_and_train
-        from ..train.lora_mlx.make_data import make_data
-
-        if not skip_preprocessing:
-            try:
-                make_data(data_dir=data_path)
-            except FileNotFoundError as exc:
-                click.secho(
-                    f"Could not read from data directory: {exc}",
-                    fg="red",
-                )
-                raise click.exceptions.Exit(1)
-
-        # NOTE we can skip this if we have a way ship MLX
-        # PyTorch safetensors to MLX safetensors
-        model_dir_local = model_path.replace("/", "-")
-        model_dir_local = f"{ckpt_output_dir}/{model_dir_local}"
-        model_dir_mlx = f"{model_dir_local}-mlx"
-        model_dir_mlx_quantized = f"{model_dir_local}-mlx-q"
-
-        if skip_quantize:
-            dest_model_dir = model_dir_mlx
-            quantize_arg = False
-        else:
-            dest_model_dir = model_dir_mlx_quantized
-            quantize_arg = True
+    if (
+        pipeline == "full"
+        or pipeline == "accelerated"
+        or strategy == SupportedTrainingStrategies.LAB_MULTIPHASE.value
+    ):
+        if not os.path.isfile(data_path):
+            ctx.fail(
+                f"Data path must be to a valid .jsonl file. Value given: {data_path}"
+            )
 
-        if tokenizer_dir is not None and gguf_model_path is not None:
-            if not local:
-                tokenizer_dir_local = tokenizer_dir.replace("/", "-")
-                fetch_tokenizer_from_hub(tokenizer_dir, tokenizer_dir_local)
+        # pull the trainrandom.randinting and torch args from the flags
+        # the flags are populated from the config as a base.
+        train_args, torch_args = map_train_to_library(ctx, ctx.params)
+        logger.debug(
+            "Rendered training arguments:\n%s", pprint.pformat(train_args.model_dump())
+        )
 
-            # no need to pass quantize_arg for now, script automatically detects if quantization is necessary based on whether gguf model is quantized or not
-            load(gguf=gguf_model_path, repo=tokenizer_dir, mlx_path=dest_model_dir)
+    if strategy == SupportedTrainingStrategies.LAB_MULTIPHASE.value:
+        if not (phased_phase1_data and phased_phase2_data):
+            raise ctx.fail(
+                "End-to-end training minimally requires: `--phased-phase1-data`, and `--phased-phase2-data`. One or more wasn't correctly specified."
+            )
 
-            for filename in os.listdir(model_dir_local):
-                shutil.copy(
-                    os.path.join(model_dir_local, filename),
-                    os.path.join(dest_model_dir, filename),
-                )
-            shutil.rmtree(model_dir_local, ignore_errors=True)
+        # if they both exist, must be Path objects
+        if not (phased_phase1_data.exists() and phased_phase2_data.exists()):
+            raise FileNotFoundError(
+                "Data for both phase1 and phase2 must be specified for phased training."
+            )
 
+        mt_bench_judge: pathlib.Path
+        if phased_mt_bench_judge is None:
+            ctx.fail("No MT-Bench model was provided with '--phased-mt-bench-judge'")
+        elif not phased_mt_bench_judge.resolve().exists():
+            raise FileNotFoundError(
+                f"MT-Bench model directory could not be found: {phased_mt_bench_judge}\nMust be an absolute path to a model directory."
+            )
         else:
-            # Downloading PyTorch SafeTensor and Converting to MLX SafeTensor
-            convert_between_mlx_and_pytorch(
-                hf_path=model_path,
-                mlx_path=dest_model_dir,
-                quantize=quantize_arg,
-                local=local,
+            # makes MyPy happy because 'mt_bench_judge' isn't Path | None
+            mt_bench_judge = phased_mt_bench_judge
+        if not skip_user_confirm:
+            click.confirm(
+                f"""\n\nSTARTING END-TO-END TRAINING.\nProceeding will OVERWRITE any metadata (checkpoints, intermediate results) that are stored at {phased_base_dir}.\n\nWould you like to proceed?""",
+                abort=True,
             )
 
-        adapter_file_path = f"{dest_model_dir}/adapters.npz"
-
-        # train the model with LoRA
-        load_and_train(
-            model=dest_model_dir,
-            train=True,
-            data=data_path,
-            adapter_file=adapter_file_path,
-            iters=iters,
-            save_every=10,
-            steps_per_eval=10,
-        )
-    elif legacy:
-        # Local
-        from ..llamacpp.llamacpp_convert_to_gguf import convert_llama_to_gguf
-        from ..train.linux_train import linux_train
+        _prepare_phased_base_dir(phased_base_dir)
 
-        training_results_dir = linux_train(
+        _run_phased_training(
             ctx=ctx,
-            train_file=train_file,
-            test_file=test_file,
-            model_name=model_path,
-            num_epochs=num_epochs,
-            train_device=device,
-            four_bit_quant=four_bit_quant,
+            train_args=train_args,
+            torch_args=torch_args,  # pylint: disable=possibly-used-before-assignment
+            base_dir=phased_base_dir,
+            phase1_data=phased_phase1_data,
+            phase1_num_epochs=phased_phase1_num_epochs,
+            phase1_samples_per_save=phased_phase1_samples_per_save,
+            phase1_checkpoints_dir=phased_base_dir / "phase1" / "checkpoints",
+            phased_phase1_effective_batch_size=phased_phase1_effective_batch_size,
+            phase2_data=phased_phase2_data,
+            phase2_num_epochs=phased_phase2_num_epochs,
+            phase2_samples_per_save=phased_phase2_samples_per_save,
+            phase2_checkpoints_dir=phased_base_dir / "phase2" / "checkpoints",
+            phased_phase2_effective_batch_size=phased_phase2_effective_batch_size,
+            phase2_eval_cache=phased_base_dir / "phase2" / "eval_cache",
+            mtbench_judge=mt_bench_judge,
+            enable_serving_output=enable_serving_output,
         )
 
-        final_results_dir = training_results_dir / "final"
-        if final_results_dir.exists():
-            shutil.rmtree(final_results_dir)
-        final_results_dir.mkdir()
-
-        gguf_models_dir = Path(DEFAULTS.CHECKPOINTS_DIR)
-        gguf_models_dir.mkdir(exist_ok=True)
-        gguf_models_file = gguf_models_dir / "ggml-model-f16.gguf"
-
-        # Remove previously trained model, its taking up space we may need in the next step
-        gguf_models_file.unlink(missing_ok=True)
-
-        # TODO: Figure out what to do when there are multiple checkpoint dirs.
-        # Right now it's just copying files from the first one numerically not necessarily the best one
-        for fpath in (
-            "checkpoint-*/added_tokens.json",
-            "checkpoint-*/special_tokens_map.json",
-            "checkpoint-*/tokenizer.json",
-            "checkpoint-*/tokenizer.model",
-            "checkpoint-*/tokenizer_config.json",
-            "merged_model/config.json",
-            "merged_model/generation_config.json",
-        ):
-            file_ = next(training_results_dir.glob(fpath))
-            shutil.copy(file_, final_results_dir)
-            print(f"Copied {file_} to {final_results_dir}")
-
-        for file in training_results_dir.glob("merged_model/*.safetensors"):
-            shutil.move(file, final_results_dir)
-            print(f"Moved {file} to {final_results_dir}")
-
-        if four_bit_quant:
-            print(
-                "SKIPPING CONVERSION to gguf. This is unsupported with --4-bit-quant. "
-                + "See https://github.com/instructlab/instructlab/issues/579."
-            )
-            return
-
-        gguf_file_path = convert_llama_to_gguf(model=final_results_dir, pad_vocab=True)
-
-        # Remove safetensors files to save space, were done with them here
-        # and the huggingface lib has them cached
-        for file in final_results_dir.glob("*.safetensors"):
-            file.unlink()
-
-        shutil.move(gguf_file_path, gguf_models_file)
-        print(f"Save trained model to {gguf_models_file}")
-
-        # cleanup checkpoint dir since it's name is unpredictable
-        # TODO: figure out how checkpoint dirs should be cleaned up
-        # checkpoint_dirs = training_results_dir.glob("checkpoint*")
-        # shutil.rmtree(checkpoint_dirs[0])
     else:
-        # run_training is a dynamic attribute, pylint is not clever enough
-        # to detect it.
-        # Third Party
-        from instructlab.training import (  # pylint: disable=no-name-in-module
-            run_training,
-        )
-
-        # pull the trainrandom.randinting and torch args from the flags
-        # the flags are populated from the config as a base.
-        train_args, torch_args = map_train_to_library(ctx, ctx.params)
-        logger.debug(
-            "Rendered training arguments:\n%s", pprint.pformat(train_args.model_dump())
-        )
+        # we can use train_args locally to run lower fidelity training
+        if is_high_fidelity(device=device) and pipeline == "accelerated":
+            # run_training is a dynamic attribute, pylint is not clever enough
+            # to detect it.
+            # Third Party
+            from instructlab.training import (
+                run_training,  # pylint: disable=no-name-in-module
+            )
 
-        if strategy == SupportedTrainingStrategies.LAB_MULTIPHASE.value:
-            if not (phased_phase1_data and phased_phase2_data):
-                raise ctx.fail(
-                    "End-to-end training minimally requires: `--phased-phase1-data`, and `--phased-phase2-data`. One or more wasn't correctly specified."
+            run_training(train_args=train_args, torch_args=torch_args)
+        elif not is_high_fidelity(device=device) and pipeline == "full":
+            # Third Party
+            import torch
+
+            torch.set_autocast_enabled(False)
+            # Local
+            from . import full_train
+
+            # if on CPU or MPS, execute full train, which is based
+            # off of the structure of the training repo, just with different optimizers, model sizes, and special data gradient accumulation to get it
+            # to fit on most consumer laptops
+            full_train.train(train_args=train_args, device=device)
+        elif pipeline == "simple":
+            if utils.is_macos_with_m_chip() and not strategy:
+                # Local
+                from ..mlx_explore.gguf_convert_to_mlx import load
+                from ..mlx_explore.utils import fetch_tokenizer_from_hub
+                from ..train.lora_mlx.convert import convert_between_mlx_and_pytorch
+                from ..train.lora_mlx.lora import load_and_train
+                from ..train.lora_mlx.make_data import make_data
+
+                if not skip_preprocessing:
+                    try:
+                        make_data(data_dir=data_path)
+                    except FileNotFoundError as exc:
+                        click.secho(
+                            f"Could not read from data directory: {exc}",
+                            fg="red",
+                        )
+                        raise click.exceptions.Exit(1)
+
+                # NOTE we can skip this if we have a way to ship MLX
+                # PyTorch safetensors to MLX safetensors
+                model_dir_local = model_path.replace("/", "-")
+                model_dir_local = f"{ckpt_output_dir}/{model_dir_local}"
+                model_dir_mlx = f"{model_dir_local}-mlx"
+                model_dir_mlx_quantized = f"{model_dir_local}-mlx-q"
+
+                if skip_quantize:
+                    dest_model_dir = model_dir_mlx
+                    quantize_arg = False
+                else:
+                    dest_model_dir = model_dir_mlx_quantized
+                    quantize_arg = True
+
+                if tokenizer_dir is not None and gguf_model_path is not None:
+                    if not local:
+                        tokenizer_dir_local = tokenizer_dir.replace("/", "-")
+                        fetch_tokenizer_from_hub(tokenizer_dir, tokenizer_dir_local)
+
+                    # no need to pass quantize_arg for now, script automatically detects if quantization is necessary based on whether gguf model is quantized or not
+                    load(
+                        gguf=gguf_model_path,
+                        repo=tokenizer_dir,
+                        mlx_path=dest_model_dir,
+                    )
+
+                    for filename in os.listdir(model_dir_local):
+                        shutil.copy(
+                            os.path.join(model_dir_local, filename),
+                            os.path.join(dest_model_dir, filename),
+                        )
+                    shutil.rmtree(model_dir_local, ignore_errors=True)
+
+                else:
+                    # Downloading PyTorch SafeTensor and Converting to MLX SafeTensor
+                    convert_between_mlx_and_pytorch(
+                        hf_path=model_path,
+                        mlx_path=dest_model_dir,
+                        quantize=quantize_arg,
+                        local=local,
+                    )
+
+                adapter_file_path = f"{dest_model_dir}/adapters.npz"
+
+                # train the model with LoRA
+                load_and_train(
+                    model=dest_model_dir,
+                    train=True,
+                    data=data_path,
+                    adapter_file=adapter_file_path,
+                    iters=iters,
+                    save_every=10,
+                    steps_per_eval=10,
                 )
-
-            # if they both exist, must be Path objects
-            if not (phased_phase1_data.exists() and phased_phase2_data.exists()):
-                raise FileNotFoundError(
-                    "Data for both phase1 and phase2 must be specified for phased training."
+            else:
+                # Local
+                from ..llamacpp.llamacpp_convert_to_gguf import convert_llama_to_gguf
+                from ..train.linux_train import linux_train
+
+                training_results_dir = linux_train(
+                    ctx=ctx,
+                    train_file=train_file,
+                    test_file=test_file,
+                    model_name=model_path,
+                    num_epochs=num_epochs,
+                    train_device=device,
+                    four_bit_quant=four_bit_quant,
                 )
 
-            mt_bench_judge: pathlib.Path
-            if phased_mt_bench_judge is None:
-                ctx.fail(
-                    "No MT-Bench model was provided with '--phased-mt-bench-judge'"
+                final_results_dir = training_results_dir / "final"
+                if final_results_dir.exists():
+                    shutil.rmtree(final_results_dir)
+                final_results_dir.mkdir()
+
+                gguf_models_dir = Path(DEFAULTS.CHECKPOINTS_DIR)
+                gguf_models_dir.mkdir(exist_ok=True)
+                gguf_models_file = gguf_models_dir / "ggml-model-f16.gguf"
+
+                # Remove previously trained model, its taking up space we may need in the next step
+                gguf_models_file.unlink(missing_ok=True)
+
+                # TODO: Figure out what to do when there are multiple checkpoint dirs.
+                # Right now it's just copying files from the first one numerically not necessarily the best one
+                for fpath in (
+                    "checkpoint-*/added_tokens.json",
+                    "checkpoint-*/special_tokens_map.json",
+                    "checkpoint-*/tokenizer.json",
+                    "checkpoint-*/tokenizer.model",
+                    "checkpoint-*/tokenizer_config.json",
+                    "merged_model/config.json",
+                    "merged_model/generation_config.json",
+                ):
+                    file_ = next(training_results_dir.glob(fpath))
+                    shutil.copy(file_, final_results_dir)
+                    logger.info(f"Copied {file_} to {final_results_dir}")
+
+                for file in training_results_dir.glob("merged_model/*.safetensors"):
+                    shutil.move(file, final_results_dir)
+                    logger.info(f"Moved {file} to {final_results_dir}")
+
+                if four_bit_quant:
+                    logger.info(
+                        "SKIPPING CONVERSION to gguf. This is unsupported with --4-bit-quant. "
+                        + "See https://github.com/instructlab/instructlab/issues/579."
+                    )
+                    return
+
+                gguf_file_path = convert_llama_to_gguf(
+                    model=final_results_dir, pad_vocab=True
                 )
-            elif not phased_mt_bench_judge.resolve().exists():
-                raise FileNotFoundError(
-                    f"MT-Bench model directory could not be found: {phased_mt_bench_judge}\nMust be an absolute path to a model directory."
-                )
-            else:
-                # makes MyPy happy because 'mt_bench_judge' isn't Path | None
-                mt_bench_judge = phased_mt_bench_judge
 
-            if not skip_user_confirm:
-                click.confirm(
-                    f"""\n\nSTARTING END-TO-END TRAINING.\nProceeding will OVERWRITE any metadata (checkpoints, intermediate results) that are stored at {phased_base_dir}.\n\nWould you like to proceed?""",
-                    abort=True,
-                )
+                # Remove safetensors files to save space, were done with them here
+                # and the huggingface lib has them cached
+                for file in final_results_dir.glob("*.safetensors"):
+                    file.unlink()
 
-            _prepare_phased_base_dir(phased_base_dir)
-
-            _run_phased_training(
-                ctx=ctx,
-                train_args=train_args,
-                torch_args=torch_args,
-                base_dir=phased_base_dir,
-                phase1_data=phased_phase1_data,
-                phase1_num_epochs=phased_phase1_num_epochs,
-                phase1_samples_per_save=phased_phase1_samples_per_save,
-                phase1_checkpoints_dir=phased_base_dir / "phase1" / "checkpoints",
-                phased_phase1_effective_batch_size=phased_phase1_effective_batch_size,
-                phase2_data=phased_phase2_data,
-                phase2_num_epochs=phased_phase2_num_epochs,
-                phase2_samples_per_save=phased_phase2_samples_per_save,
-                phase2_checkpoints_dir=phased_base_dir / "phase2" / "checkpoints",
-                phased_phase2_effective_batch_size=phased_phase2_effective_batch_size,
-                phase2_eval_cache=phased_base_dir / "phase2" / "eval_cache",
-                mtbench_judge=mt_bench_judge,
-                enable_serving_output=enable_serving_output,
-            )
+                shutil.move(gguf_file_path, gguf_models_file)
+                logger.info(f"Save trained model to {gguf_models_file}")
 
-        else:
-            if not os.path.isfile(data_path):
-                ctx.fail(
-                    f"Data path must be to a valid .jsonl file. Value given: {data_path}"
-                )
 
-            # TODO: should we have this wrapped in a try-except?
-            run_training(train_args=train_args, torch_args=torch_args)
+# chooses which type of training to run depending on the device provided
+def is_high_fidelity(device):
+    return device in ("cuda", "hpu")
 
 
 def _prepare_phased_base_dir(phased_base_dir: pathlib.Path) -> None:
@@ -794,7 +824,6 @@ def _mtbench(
     enable_serving_output: bool,
 ) -> float:
     # TODO: optimization: run all generations in serial and then do all judgments at once to save time loading/unloading prometheus.
-
     # Third Party
     from instructlab.eval.mt_bench import MTBenchEvaluator
     import torch
diff --git a/tox.ini b/tox.ini
index d07969f5f4..5cc8fadf95 100644
--- a/tox.ini
+++ b/tox.ini
@@ -142,6 +142,7 @@ deps =
   types-PyYAML
   types-requests
   types-tqdm
+  types-psutil
   pytest
 commands =
   mypy {posargs}


Test patch:
diff --git a/.github/workflows/e2e-nvidia-a10g-x1.yml b/.github/workflows/e2e-nvidia-a10g-x1.yml
index fe0e9e2909..6c3673ee20 100644
--- a/.github/workflows/e2e-nvidia-a10g-x1.yml
+++ b/.github/workflows/e2e-nvidia-a10g-x1.yml
@@ -31,7 +31,7 @@ jobs:
           mode: start
           github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
           ec2-image-id: ami-00c51d9c1374eda97
-          ec2-instance-type: g5.2xlarge
+          ec2-instance-type: g5.4xlarge
           subnet-id: subnet-02d230cffd9385bd4
           security-group-id: sg-06300447c4a5fbef3
           iam-role-name: instructlab-ci-runner
@@ -159,7 +159,7 @@ jobs:
 
           python3.11 -m pip show nvidia-nccl-cu12
           
-          ./scripts/basic-workflow-tests.sh -mL
+          ./scripts/basic-workflow-tests.sh -mf
 
       - name: Add comment to PR if the workflow failed
         if: failure() && steps.check_pr.outputs.is_pr == 'true'
diff --git a/.github/workflows/e2e-nvidia-a10g-x4.yml b/.github/workflows/e2e-nvidia-a10g-x4.yml
index a6f45c80ae..6e08371d6d 100644
--- a/.github/workflows/e2e-nvidia-a10g-x4.yml
+++ b/.github/workflows/e2e-nvidia-a10g-x4.yml
@@ -148,7 +148,7 @@ jobs:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
           . venv/bin/activate
-          ./scripts/basic-workflow-tests.sh -mevFMT
+          ./scripts/basic-workflow-tests.sh -mevFMa
 
       - name: Add comment to PR if the workflow failed
         if: failure() && steps.check_pr.outputs.is_pr == 'true'
diff --git a/.github/workflows/e2e-nvidia-t4-x1.yml b/.github/workflows/e2e-nvidia-t4-x1.yml
index 4e6a0b2e78..425fd2ef58 100644
--- a/.github/workflows/e2e-nvidia-t4-x1.yml
+++ b/.github/workflows/e2e-nvidia-t4-x1.yml
@@ -121,7 +121,7 @@ jobs:
       - name: Run e2e test
         run: |
           . venv/bin/activate
-          ./scripts/basic-workflow-tests.sh -m
+          ./scripts/basic-workflow-tests.sh -msq
 
   stop-runner:
     name: Stop external EC2 runner
diff --git a/scripts/basic-workflow-tests.sh b/scripts/basic-workflow-tests.sh
index f928824b58..11a9392990 100755
--- a/scripts/basic-workflow-tests.sh
+++ b/scripts/basic-workflow-tests.sh
@@ -17,10 +17,12 @@ NUM_INSTRUCTIONS=5
 GENERATE_ARGS=("--num-cpus" "$(nproc)" --taxonomy-path='./taxonomy')
 DIFF_ARGS=("--taxonomy-path" "./taxonomy")
 TRAIN_ARGS=()
-LEGACYTRAIN=0
 PHASED_TRAINING=0
-TRAIN_LIBRARY=0
 BACKEND="llama-cpp"
+FOUR_BIT_QUANT=0
+SIMPLE_TRAIN=0
+FULL_TRAIN=0
+ACCELERATED_TRAIN=0
 HF_TOKEN=${HF_TOKEN:-}
 SDG_PIPELINE="simple"
 SKIP_TRAIN=${SKIP_TRAIN:-0}
@@ -97,15 +99,14 @@ set_defaults() {
         exit 1
     fi
 
-    if [ "${PHASED_TRAINING}" -eq 1 ] && [ "${TRAIN_LIBRARY}" -eq 0 ]; then
-        echo "ERROR: You have -P set. It requires -T."
+    if [ "${PHASED_TRAINING}" -eq 1 ] && [ "${ACCELERATED_TRAIN}" -eq 0 ]; then
+        echo "ERROR: You have -P set. It requires -a."
         exit 1
     fi
 
     if [ "$MINIMAL" -eq 1 ]; then
         # Minimal settings to run in less time
         NUM_INSTRUCTIONS=1
-        TRAIN_ARGS+=("--num-epochs" "1")
     fi
 }
 
@@ -116,14 +117,8 @@ test_smoke() {
 
 test_init() {
     task Initializing ilab
-
-    if [ "$LEGACYTRAIN" -eq 1 ]; then
-        # TODO Only cuda for now
-        step Setting train-profile for GPU accelerated training
-        ilab config init --non-interactive --train-profile="${SCRIPTDIR}/test-data/train-profile-a10.yaml"
-    else
-        ilab config init --non-interactive
-    fi
+    
+    ilab config init --non-interactive
 
     step Replace model in config.yaml
     if [ "${BACKEND}" == "vllm" ]; then
@@ -283,29 +278,36 @@ test_generate() {
 test_train() {
     task Train the model
 
-    if [ "$TRAIN_LIBRARY" -eq 1 ]; then
-        local data
-        data=$(find "${DATA_HOME}"/instructlab/datasets -name 'messages_*' | head -n 1)
+    local data
+    data=$(find "${DATA_HOME}"/instructlab/datasets -name 'skills_train_msgs_*' | head -n 1)
+
+    # simple, full, and accelerated, are different workflows
+    # To mimic a real user e2e scenario, only one of these should be run on a given system
+    # The `small` worker can manage `simple`, The medium worker can handle `full` and the large worker can handle `accelerated`
+    if [ "$ACCELERATED_TRAIN" -eq 1 ]; then
         # TODO Only cuda for now
         # the train profile specified in test_init overrides the majority of TRAIN_ARGS, including things like num_epochs. While it looks like much of those settings are being lost, they just have different values here.
-        TRAIN_ARGS=("--device=cuda" "--model-path=${GRANITE_SAFETENSOR_REPO}" "--data-path=${data}" "--lora-quantize-dtype=nf4" "--4-bit-quant" "--effective-batch-size=4" "--is-padding-free=False")
+        TRAIN_ARGS=("--pipeline=accelerated" "--device=cuda" "--model-path=${GRANITE_SAFETENSOR_REPO}" "--data-path=${data}" "--lora-quantize-dtype=nf4" "--4-bit-quant" "--effective-batch-size=4" "--is-padding-free=False")
         if [ "${BACKEND}" != "vllm" ]; then
             TRAIN_ARGS+=("--gguf-model-path" "${CACHE_HOME}/instructlab/models/${GRANITE_GGUF_MODEL}")
         fi
-
-        ilab model train "${TRAIN_ARGS[@]}"
-    else
-        # TODO Only cuda for now
-        TRAIN_ARGS+=("--legacy" "--device=cuda")
-        if [ "$LEGACYTRAIN" -eq 0 ]; then
+    fi
+    if [ "$SIMPLE_TRAIN" -eq 1 ]; then
+        if [ "$FOUR_BIT_QUANT" -eq 1 ]; then
             TRAIN_ARGS+=("--4-bit-quant")
         fi
+        # TODO Only cuda for now
+        TRAIN_ARGS+=("--pipeline=simple" "--device=cuda" "--num-epochs=1")
         if [ "${BACKEND}" != "vllm" ]; then
             TRAIN_ARGS+=("--gguf-model-path" "${CACHE_HOME}/instructlab/models/${GRANITE_GGUF_MODEL}")
         fi
-
-        ilab model train "${TRAIN_ARGS[@]}"
     fi
+    if [ "$FULL_TRAIN" -eq 1 ]; then
+        # test training on a CPU not the GPU
+        TRAIN_ARGS=("--num-epochs=1" "--pipeline=full" "--model-path=${GRANITE_SAFETENSOR_REPO}" "--data-path=${data}" "--effective-batch-size=4" --device=cpu)
+    fi
+
+    ilab model train "${TRAIN_ARGS[@]}"
 }
 
 test_phased_train() {
@@ -451,15 +453,18 @@ test_exec() {
     # When we run training with --4-bit-quant, we can't convert the result to a gguf
     # https://github.com/instructlab/instructlab/issues/579
     # so we skip trying to test the result
-    if [ "$LEGACYTRAIN" -eq 1 ]; then
+    if [ "$FULL_TRAIN" -eq 1 ]; then
         # When you run this --
         #   `ilab model convert` is only implemented for macOS with M-series chips for now
         #test_convert
-    
-        test_serve trained "${DATA_HOME}/instructlab/checkpoints/model.gguf"
+
+        # when using full train, choose any GGUF from any of the checkpoints dirs
+        model_dir=$(find "${DATA_HOME}"/instructlab/checkpoints/hf_format -name 'samples_*' | head -n 1)
+
+        test_serve trained "${model_dir}/pytorch_model-Q4_K_M.gguf"
         PID=$!
 
-        test_chat
+        ilab model chat -qq --model "${model_dir}/pytorch_model-Q4_K_M.gguf" --endpoint-url http://localhost:8000/v1 'Say "Hello" and nothing else\n'
 
         # Kill the serve process
         task Stopping the ilab model serve for trained model
@@ -508,21 +513,21 @@ wait_for_server() {
 usage() {
     echo "Usage: $0 [-m] [-h]"
     echo "  -e  Run model evaluation"
-    echo "  -T  Use the 'full' training library rather than legacy training"
+    echo "  -q  Use 4-bit-quant when training"
+    echo "  -a  Use the 'full' training library rather than legacy training"
+    echo "  -s  Run the simple training using the SFTTainer rather than the custom training loop"
     echo "  -f  Run the fullsize training instead of --4-bit-quant"
     echo "  -F  Use the 'full' SDG pipeline instead of the default 'simple' pipeline"
     echo "  -h  Show this help text"
-    echo "  -L  Run legacy training with 4-bit quantization"
     echo "  -m  Run minimal configuration with lower number of instructions and training epochs (run quicker when you have no GPU)"
     echo "  -M  Use the mixtral model (4-bit quantized) instead of merlinite model (4-bit quantized)."
     echo "  -P  Run multi-phase training"
-    echo "  -T  Use the 'full' training library rather than legacy training"
     echo "  -v  Use the vLLM backend for serving"
 }
 
 # Process command line arguments
 task "Configuring ..."
-while getopts "eFhLmMPTv" opt; do
+while getopts "eFhqasfmMPv" opt; do
     case $opt in
         e)
             EVAL=1
@@ -536,10 +541,6 @@ while getopts "eFhLmMPTv" opt; do
             usage
             exit 0
             ;;
-        L)
-            LEGACYTRAIN=1
-            step "Running legacy training with 4-bit quantization."
-            ;;
         m)
             MINIMAL=1
             step "Running minimal configuration."
@@ -552,14 +553,26 @@ while getopts "eFhLmMPTv" opt; do
             PHASED_TRAINING=1
             step "Running multi-phase training."
             ;;
-        T)
-            TRAIN_LIBRARY=1
-            step "Running with training library."
-            ;;
         v)
             BACKEND=vllm
             step "Running with vLLM backend."
             ;;
+        q)
+            FOUR_BIT_QUANT=1
+            step "Running training using 4-bit-quantization."
+            ;;
+        s)
+            SIMPLE_TRAIN=1
+            step "Running the simple training pipeline"
+            ;;
+        f)
+            FULL_TRAIN=1
+            step "Running the full training pipeline"
+            ;;
+        a)
+            ACCELERATED_TRAIN=1
+            step "Running using the training library"
+            ;;
         \?)
             echo "Invalid option: -$opt" >&2
             usage
diff --git a/tests/test_lab_train.py b/tests/test_lab_train.py
index 48c399cc1f..78d7479677 100644
--- a/tests/test_lab_train.py
+++ b/tests/test_lab_train.py
@@ -94,7 +94,15 @@ def test_train_mac(
         setup_input_dir()
         result = cli_runner.invoke(
             lab.ilab,
-            ["--config=DEFAULT", "model", "train", "--input-dir", INPUT_DIR],
+            [
+                "--config=DEFAULT",
+                "model",
+                "train",
+                "--pipeline",
+                "simple",
+                "--input-dir",
+                INPUT_DIR,
+            ],
         )
         assert result.exit_code == 0
         load_mock.assert_not_called()
@@ -139,6 +147,8 @@ def test_skip_quantize(
                 "--config=DEFAULT",
                 "model",
                 "train",
+                "--pipeline",
+                "simple",
                 "--input-dir",
                 INPUT_DIR,
                 "--skip-quantize",
@@ -159,7 +169,8 @@ def test_input_error(self, cli_runner: CliRunner):
                 "--config=DEFAULT",
                 "model",
                 "train",
-                "--legacy",
+                "--pipeline",
+                "simple",
                 "--input-dir",
                 "invalid",
             ],
@@ -176,7 +187,8 @@ def test_invalid_taxonomy(self, cli_runner: CliRunner):
                 "--config=DEFAULT",
                 "model",
                 "train",
-                "--legacy",
+                "--pipeline",
+                "simple",
                 "--input-dir",
                 INPUT_DIR,
             ],
@@ -198,7 +210,8 @@ def test_invalid_data_dir(self, cli_runner: CliRunner):
                     "--config=DEFAULT",
                     "model",
                     "train",
-                    "--legacy",
+                    "--pipeline",
+                    "simple",
                     "--data-path",
                     "invalid",
                     "--input-dir",
@@ -224,7 +237,8 @@ def test_invalid_data_dir_synthetic(
                 "--config=DEFAULT",
                 "model",
                 "train",
-                "--legacy",
+                "--pipeline",
+                "simple",
                 "--data-path",
                 "invalid",
                 "--input-dir",
@@ -258,6 +272,8 @@ def test_skip_preprocessing(
                 "--config=DEFAULT",
                 "model",
                 "train",
+                "--pipeline",
+                "simple",
                 "--input-dir",
                 INPUT_DIR,
                 "--skip-preprocessing",
@@ -294,7 +310,8 @@ def test_load(
                 "--config=DEFAULT",
                 "model",
                 "train",
-                "--legacy",
+                "--pipeline",
+                "simple",
                 "--input-dir",
                 DEFAULTS.DATASETS_DIR,
                 "--tokenizer-dir",
@@ -340,7 +357,8 @@ def test_load_local(
                 "--config=DEFAULT",
                 "model",
                 "train",
-                "--legacy",
+                "--pipeline",
+                "simple",
                 "--input-dir",
                 DEFAULTS.DATASETS_DIR,
                 "--tokenizer-dir",
@@ -382,7 +400,8 @@ def test_train_linux(
                 "--config=DEFAULT",
                 "model",
                 "train",
-                "--legacy",
+                "--pipeline",
+                "simple",
                 "--input-dir",
                 INPUT_DIR,
             ],
@@ -450,7 +469,8 @@ def test_double_train_linux(
                     "--config=DEFAULT",
                     "model",
                     "train",
-                    "--legacy",
+                    "--pipeline",
+                    "simple",
                     "--input-dir",
                     DEFAULTS.DATASETS_DIR,
                     "--data-path",
@@ -489,7 +509,8 @@ def test_double_train_linux(
                     "--config=DEFAULT",
                     "model",
                     "train",
-                    "--legacy",
+                    "--pipeline",
+                    "simple",
                     "--input-dir",
                     DEFAULTS.DATASETS_DIR,
                     "--data-path",
@@ -540,7 +561,8 @@ def test_num_epochs(
                 "--config=DEFAULT",
                 "model",
                 "train",
-                "--legacy",
+                "--pipeline",
+                "simple",
                 "--input-dir",
                 INPUT_DIR,
                 "--num-epochs",