Commit code to repositority for initial basic training runs

[vishalbakshi]

Dependent on #11 and #9.

Do as much of our testing on SmolLM2 135M/400M models before we scale up.

• Make sure rsLoRA, LoRA, DoRA works on the models we pick and we have the command to run it all
• Perform rigorous hyperparameter sweep to understand differences betwen rsLoRA, LoRA and DoRA. Don't have to be full runs and so can be done parallel with faster training pipelines being implemented.
• We need tests that our LoRA/rsLoRA/DoRA training result matches torchtune/HF/etc.

[vishalbakshi]

I'm currently working on documenting quick starts using SmolLM2-135M in the docs branch.

[matdmiller]

Here is a first step towards getting a modal script set up to do real training runs:

• Add persistence for datasets and model checkpoints with volumes.
• Started to create a more consistent folder structure for generated data.
• Add ability to view aim training run stats live. Look for printed line just before training starts that looks like this Aim server available at: https://k4mrv1mslxmw81.r18.modal.host. This link will work while the training run is active and then stop working once it completes after which time you need to look at the run in the centralized repo. This will be handy for long training runs.
• Split out separate steps into different functions that can be run independently.
• Store state including datasets, model checkpoints, converted hf models, eval stats

@vishalbakshi @galopyz

# INSTRUCTIONS
# Create volumes (one time via cli) - run: `modal volume create lrg-datasets` and `modal volume create lrg-model-checkpoints`
#TODO:
# - Probably want to put local aim repo in the folder with the model checkpoints
# - Fix duplicate container issue


import modal
from modal import Image, App, Secret, Volume

import pathlib, datetime

PYTHON_PATH = "/opt/conda/envs/llm-foundry/bin/python"
TRAINING_GPU = "h100" # "a10g" "h100" # "l4"
BATCH_SIZE = 20 # 20 for h100 4 for l4
DATASET_BASE_PATH = "/datasets"
DATASETS_VOLUME = Volume.from_name("lrg-datasets")
DATASETS_VOLUME_MOUNT_PATH = pathlib.Path("/datasets")
MODEL_CHECKPOINT_VOLUME = Volume.from_name("lrg-model-checkpoints")
MODEL_CHECKPOINT_VOLUME_MOUNT_PATH = pathlib.Path("/model-checkpoints")

app = App("quick-start")

# Build image from local Dockerfile
image = Image.from_dockerfile("Dockerfile")

@app.function(gpu=TRAINING_GPU, image=image, timeout=3600, secrets=[Secret.from_name("LRG")],
             concurrency_limit=1)
def get_stats():
    import subprocess
    
    # Use the correct Python interpreter for imports
    import_check = subprocess.run(
        [PYTHON_PATH, "-c", "import flash_attn; print(flash_attn.__version__)"],
        capture_output=True,
        text=True
    )
    print(f"Flash Attention version: {import_check.stdout}")

    # Run nvidia-smi to check GPU status
    nvidia_smi = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
    print("NVIDIA-SMI Output:")
    print(nvidia_smi.stdout)
    if nvidia_smi.stderr:
        print("NVIDIA-SMI Errors:", nvidia_smi.stderr)


@app.function(gpu=TRAINING_GPU, image=image, timeout=3600, secrets=[Secret.from_name("LRG")], 
              volumes={DATASETS_VOLUME_MOUNT_PATH: DATASETS_VOLUME},
              concurrency_limit=1)
def convert_c4_small_dataset():
    import subprocess
    import os
    
    # Change to llm-foundry/scripts directory at the start
    os.chdir("/llm-foundry/scripts")
    print(f"Working directory: {os.getcwd()}")
    
    # Step 1: Convert C4 dataset
    print("Converting C4 dataset...")
    data_prep_cmd = [
        PYTHON_PATH,  # Use the correct Python interpreter
        "data_prep/convert_dataset_hf.py",
        "--dataset", "allenai/c4",
        "--data_subset", "en",
        "--out_root", f"{DATASET_BASE_PATH}/c4_small",
        "--splits", "train_small", "val_small",
        "--concat_tokens", "2048",
        "--tokenizer", "HuggingFaceTB/SmolLM2-135M",
        "--eos_text", "<|endoftext|>"
    ]
    result = subprocess.run(data_prep_cmd, capture_output=True, text=True)
    print(result.stdout)
    if result.stderr:
        print("Data prep errors:", result.stderr)
    
    DATASETS_VOLUME.commit()

@app.function(gpu=TRAINING_GPU, image=image, timeout=3600, secrets=[Secret.from_name("LRG")], 
              volumes={MODEL_CHECKPOINT_VOLUME_MOUNT_PATH: MODEL_CHECKPOINT_VOLUME},
              concurrency_limit=1)
def view_model_checkpoints(save_folder: str=None):
    import os
    print("\nModel checkpoint files and sizes:")
    for folder in ([save_folder] if save_folder else os.listdir(MODEL_CHECKPOINT_VOLUME_MOUNT_PATH)):
        folder = os.path.join(MODEL_CHECKPOINT_VOLUME_MOUNT_PATH, folder)
        if os.path.isdir(folder):
            for filename in os.listdir(folder):
                filepath = os.path.join(folder, filename)
                if os.path.isfile(filepath):
                    size_mb = os.path.getsize(filepath) / (1024 * 1024)
                    print(f"{os.path.join(folder, filename)}: {size_mb:.2f} MB")


# @app.function(gpu=TRAINING_GPU, image=image, timeout=6*3600, secrets=[Secret.from_name("LRG")], 
#               volumes={MODEL_CHECKPOINT_VOLUME_MOUNT_PATH: MODEL_CHECKPOINT_VOLUME,
#                        DATASETS_VOLUME_MOUNT_PATH: DATASETS_VOLUME})
def train_model(run_ts: str):
    import subprocess
    import os
    
    # Change to llm-foundry/scripts directory at the start
    os.chdir("/llm-foundry/scripts")
    print(f"Working directory: {os.getcwd()}")
    
    # Step 2: Train the model
    print("\nTraining model...")
    save_folder = f"{MODEL_CHECKPOINT_VOLUME_MOUNT_PATH}/smollm2-135m-{run_ts}"
    
    train_cmd = [
        "composer",
        "train/train.py",
        "train/yamls/pretrain/smollm2-135m.yaml",  # Updated YAML path
        "loggers.aim.experiment_name=quickstart_test_smollm2_135m_modal",
        f"variables.data_local={DATASETS_VOLUME_MOUNT_PATH}/c4_small",
        "train_loader.dataset.split=train_small",
        "eval_loader.dataset.split=val_small",
        "max_duration=500ba",
        "eval_interval=500ba",
        f"save_folder={save_folder}",  # Updated model name
        f"device_eval_batch_size={BATCH_SIZE}",  # Added batch size settings # 20 for h100 4 for l4
        f"device_train_microbatch_size={BATCH_SIZE}",
        f"global_train_batch_size={BATCH_SIZE}"
    ]
    result = subprocess.run(train_cmd, capture_output=True, text=True)
    print(result.stdout)
    print(f'Training complete for {run_ts}')
    print(f'Model checkpoints saved to {save_folder}')

    MODEL_CHECKPOINT_VOLUME.commit()

    # Print checkpoint file sizes
    view_model_checkpoints.remote(save_folder)
    
    if result.stderr:
        print("Training errors:", result.stderr)
    if result.returncode != 0:
        raise Exception(f"Training failed with exit code {result.returncode}\nStderr: {result.stderr}")
    return str(save_folder)

def run_aim_server():
    import subprocess
    import os
    
    os.chdir("/llm-foundry/scripts")
    print("Initializing Aim...")
    subprocess.run(["aim", "init"], check=True)
    
    process = subprocess.Popen(
        ["aim", "up", "--host", "0.0.0.0", "--port", "43800"],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )
    
    return process


@app.function(gpu=TRAINING_GPU, image=image, timeout=12*3600, secrets=[Secret.from_name("LRG")],
              volumes={MODEL_CHECKPOINT_VOLUME_MOUNT_PATH: MODEL_CHECKPOINT_VOLUME,
                      DATASETS_VOLUME_MOUNT_PATH: DATASETS_VOLUME},
              concurrency_limit=1)
def train_with_aim(run_ts: str):
    import subprocess, time

    with modal.forward(43800) as tunnel:
        print(f"\nAim server available at: {tunnel.url}")
        model_path = None
        aim_task = run_aim_server()
        time.sleep(5)
    
        try:
            model_path = train_model(run_ts)

        finally:
            aim_task.terminate()
            try:
                aim_task.wait(timeout=5)
            except subprocess.TimeoutExpired:
                aim_task.kill()
    
    return model_path

@app.function(gpu=TRAINING_GPU, image=image, timeout=3600, secrets=[Secret.from_name("LRG")],
              volumes={MODEL_CHECKPOINT_VOLUME_MOUNT_PATH: MODEL_CHECKPOINT_VOLUME},
              concurrency_limit=1)
def convert_model_to_hf(checkpoint_path: str):
    """
    Convert a model checkpoint to a HuggingFace format.
    """
    import subprocess, os
    from pathlib import Path

    os.chdir("/llm-foundry/scripts")
    print(f"Working directory: {os.getcwd()}")

    composer_checkpoint_path = Path(MODEL_CHECKPOINT_VOLUME_MOUNT_PATH)/checkpoint_path
    if composer_checkpoint_path.is_dir():
        composer_checkpoint_path = composer_checkpoint_path / "latest-rank0.pt"
    hf_output_path = composer_checkpoint_path.parent / "hf"

    print("\nConverting model to HuggingFace format...")
    convert_cmd = [
        PYTHON_PATH, "inference/convert_composer_to_hf.py",
        "--composer_path", composer_checkpoint_path,
        "--hf_output_path", hf_output_path,
        "--output_precision", "bf16"
    ]
    result = subprocess.run(convert_cmd, capture_output=True, text=True)
    print(result.stdout)
    if result.stderr:
        print("Conversion errors:", result.stderr)
    MODEL_CHECKPOINT_VOLUME.commit()
    print("Conversion complete!")


@app.function(gpu=TRAINING_GPU, image=image, timeout=3600, secrets=[Secret.from_name("LRG")],
              volumes={MODEL_CHECKPOINT_VOLUME_MOUNT_PATH: MODEL_CHECKPOINT_VOLUME},
              concurrency_limit=1)
def evaluate_model(checkpoint_path: str):
    import subprocess, os
    from pathlib import Path

    os.chdir("/llm-foundry/scripts")
    print(f"Working directory: {os.getcwd()}")
    
    model_path = Path(MODEL_CHECKPOINT_VOLUME_MOUNT_PATH)/checkpoint_path
    save_path = model_path.parent/"evals"  # Create evals subfolder path
    
    print("\nEvaluating model...")
    eval_cmd = [
        "composer",
        "eval/eval.py",
        "eval/yamls/hf_eval.yaml",
        "icl_tasks=eval/yamls/copa.yaml",
        f"variables.model_name_or_path={model_path}",
        f"results_path={save_path}"  # Add results_path parameter
    ]
    result = subprocess.run(eval_cmd, capture_output=True, text=True)
    print(result.stdout)
    if result.stderr:
        print("Evaluation errors:", result.stderr)
    
    MODEL_CHECKPOINT_VOLUME.commit()  # Commit the new eval results
    print("Evaluation complete!")


@app.function(gpu=TRAINING_GPU, image=image, timeout=3600, secrets=[Secret.from_name("LRG")],
              volumes={MODEL_CHECKPOINT_VOLUME_MOUNT_PATH: MODEL_CHECKPOINT_VOLUME},
              concurrency_limit=1)
def generate_responses(checkpoint_path: str, prompts: list[str]|str|None=None):
    import subprocess, os
    from pathlib import Path

    os.chdir("/llm-foundry/scripts")
    print(f"Working directory: {os.getcwd()}")
    
    model_path = Path(MODEL_CHECKPOINT_VOLUME_MOUNT_PATH)/checkpoint_path

    if prompts is None:
        prompts = [
            "The answer to life, the universe, and happiness is",
            "Here's a quick recipe for baking chocolate chip cookies: Start by"
        ]
    elif isinstance(prompts, str):
        prompts = [prompts]
    

    print("\nGenerating test responses...")
    generate_cmd = [
        PYTHON_PATH, "inference/hf_generate.py",
        "--name_or_path", model_path,
        "--max_new_tokens", "256",
        "--prompts",
        *prompts
    ]
    result = subprocess.run(generate_cmd, capture_output=True, text=True)
    print(result.stdout)
    if result.stderr:
        print("Generation errors:", result.stderr)
    print("Generation complete!")


@app.local_entrypoint()
def main():
    from pathlib import Path
    import time
    run_ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")

    get_stats.remote()
    time.sleep(1)
    # convert_c4_small_dataset.remote() # Only run once

    model_path = train_with_aim.remote(run_ts)
    print(f"Model path: {model_path}")
    model_path = Path(model_path).name
    time.sleep(1)
    
    view_model_checkpoints.remote()
    time.sleep(1)
    convert_model_to_hf.remote(model_path)
    time.sleep(1)
    evaluate_model.remote(model_path+"/hf")
    time.sleep(1)
    generate_responses.remote(model_path+"/hf")

[galopyz]

This looks awesome! So many improvements everyday.
I think this is a great rough draft.
Thank you. 🚀

[vishalbakshi]

@matdmiller This is fantastic!! I particularly am excited to see the modular nature of the code and the real time AIM UI updates (I had resigned myself to the fact that our log copying utility would not allow it).

I ran your script twice successfully and saw the live logging each time and saw all of the print outs throughout the script (nvidia-smi, checkpoints, etc.):

• L4 GPU
  • batch size = 4
  • Ran convert_c4_small_dataset.remote()
• A100-80GB GPU
  • batch size = 20
  • Did not run convert_c4_small_dataset.remote()

One question: why does it save two checkpoints each time?

First L4 run:

Model checkpoint files and sizes:
/model-checkpoints/smollm2-135m-20250212_064224/ep0-ba500-rank0.pt: 1544.14 MB
/model-checkpoints/smollm2-135m-20250212_064224/latest-rank0.pt: 1544.14 MB

Second run:

Model checkpoint files and sizes:
/model-checkpoints/smollm2-135m-20250212_064224/ep0-ba500-rank0.pt: 1544.14 MB
/model-checkpoints/smollm2-135m-20250212_064224/latest-rank0.pt: 1544.14 MB
/model-checkpoints/smollm2-135m-20250212_070525/ep0-ba600-rank0.pt: 1544.14 MB
/model-checkpoints/smollm2-135m-20250212_070525/latest-rank0.pt: 1544.14 MB

---

A few other notes:

From their pricing page:

We currently do not charge for Volumes, which are our native high-performance file serving objects. This is likely to change in early 2025. When we do start charging for Volumes, you can expect the price to be comparable to Amazon S3 pricing.

Something to keep an eye on.

---

A couple "future work" suggestions (low priority, should not keep us from merging a PR on this):

• From their docs, "You can also create Volumes lazily from code", just to remove one additional step using CLI

vol = modal.Volume.from_name("my-volume", create_if_missing=True)

• Adding/testing how to delete a checkpoint (if storage ever becomes an issue) but with a note that it's better to use some credits on storage than risk losing a checkpoint we later realize we need.
• Adding a function to upload checkpoints to HF (for final training runs).
• We might want to consider naming the run and the checkpoint the same (and perhaps add _{gpu name} at the end).

[matdmiller]

Yea, I'm really glad this live view works too. I was not satisfied with not having visibility into long training runs :) The last item in the checkpoint list is a symlink, not a separate file. In your list there is actually just 1 file and 1 symlink. By default currently it seems to checkpoint every 1,000 batches and that symlink always points to the latest checkpoint. I'm sure that is configurable but have not looked into it beyond that.

ex:

Re: Volume storage costs
I think this will be fairly minimal for us as long as we clean up things we don't care about saving. The only unknown for me is what the access costs will look like (if any) and what our usage pattern looks like. I don't think it will be significant as long as we pay at least a little attention to it.

Good to know on the auto volume creation, that is simpler.

For adding/deleting checkpoints should be as simple as writing a quick function in python to delete whatever we want. It's just python and files on a file system from python's perspective.

Agreed on adding a function to upload to HF.

The problem is that aim doesn't create the random id until you actually run the training and one of the parameters for the training command is the checkpoint name. I think we will be able to add the ability to add aim "tags" that capture the checkpoint name to make it a bit easier to match up the 2, but at least for now, the checkpoint name is captured by aim when it logs, you just have to dig a bit for it. I think we can come up with a easier to use way nonetheless.

[vishalbakshi]

Awesome, that all makes sense. And thanks for clarifying about the symlink I see that on my end as well

[galopyz]

I was able to run twice as well.
Once with and once without convert_c4_small_dataset.remote()

I used a10g GPU with 4 batch size, and H100 with 20 bs.
I confirm AIM logging worked and found the symlink as well.
Pretty cool stuff.

[matdmiller]

Created a branch with this and the start to a README. Only a couple very minor changes compared to what I pasted above. Please feel free to comment and contribute. I'll work on this some more later.

Leaving this note here on how to add more dependencies to our uv.lock file for the huggingface_hub client. (I'm new to using uv) https://docs.astral.sh/uv/guides/projects/#managing-dependencies

[vishalbakshi]

@matdmiller @galopyz documenting this here for future reference---maybe this is obvious (wasn't to me) but the validation step is the significant bottleneck during training:

I'm using an L4 GPU.

max_duration=100ba

eval_interval=0:
100 ba x 1 (bs) x 2048 (seq len) = 204800 tokens in 29 seconds = 25.4M tokens/hour

eval_interval=100ba:
100 ba x 1 (bs) x 2048 (seq len) = 204800 tokens in 8.6 minutes = 1.4M tokens/hour

max_duration=500ba

eval_interval=0:
500 ba x 1 (bs) x 2048 (seq len) = 1024000 tokens in 107 seconds = 34.5M tokens/hour

eval_interval=500ba:
500 ba x 1 (bs) x 2048 (seq len) = 1024000 tokens in 10 minutes = 6.1M tokens/hour