> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # Model Training on Databricks Databricks is a widely used platform for managing data, models, applications, and compute on the cloud. This guide shows how to use AutoModel for scalable, performant model training on Databricks. The specific example here fine-tunes a [Llama-3.2-1B](https://huggingface.co/meta-llama/Llama-3.2-1B) model using the [SQuAD dataset](https://huggingface.co/datasets/rajpurkar/squad) from Hugging Face, but any AutoModel functionality (for example, [model pre-training](/recipes-e2e-examples/pretraining), [VLMs](/model-coverage/vision-language-models/overview), [other supported models](/model-coverage/overview)) can also be run on Databricks. ## Provision Compute Start by [provisioning](https://docs.databricks.com/aws/en/compute/configure) a Databricks classic compute cluster with the following setup: * Databricks runtime: [18.0 LTS (Machine Learning version)](https://docs.databricks.com/aws/en/release-notes/runtime/18.0ml) * Worker instance type: `g6e.12xlarge` on AWS (4x L40S GPUs per node) * Number of workers: 2 * Global [environment variable](https://docs.databricks.com/aws/en/compute/configure#environment-variables): `GLOO_SOCKET_IFNAME=eth0` (see [this](https://docs.databricks.com/aws/en/machine-learning/train-model/distributed-training/spark-pytorch-distributor#gloo-failure-runtimeerror-connection-refused) for details) * Cluster-scoped [init script](https://docs.databricks.com/aws/en/init-scripts/cluster-scoped): ```bash #!/bin/bash # Install uv, then install AutoModel into the Databricks Python environment on all nodes. curl -LsSf https://astral.sh/uv/install.sh | sh /root/.local/bin/uv pip install --python /databricks/python3/bin/python \ "nemo-automodel @ git+https://github.com/NVIDIA-NeMo/Automodel.git" ``` This provisions three compute nodes: one driver node that you can attach a notebook to, and two worker nodes to use for multi-node training. Note that this guide uses a small number of instances for demonstration purposes, but you can adjust the specific instance type and number of workers for your use case. ## Train the Model With the provisioned compute resources, you are ready to fine-tune a model using AutoModel. AutoModel uses YAML file recipes to configure various settings for the training process, such as the model, dataset, loss function, and optimizer. This guide uses this [preconfigured recipe](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml) for fine-tuning a Llama-3.2-1B model using the SQuAD dataset from Hugging Face. In a notebook connected to the compute resource, download the configuration file: ```bash # Download configuration file !curl -O https://raw.githubusercontent.com/NVIDIA-NeMo/Automodel/refs/heads/main/examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml ``` The model, dataset, and optimizer portions of the configuration file are shown below: ```yaml model: _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained pretrained_model_name_or_path: meta-llama/Llama-3.2-1B dataset: _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset dataset_name: rajpurkar/squad split: train optimizer: _target_: torch.optim.Adam betas: [0.9, 0.999] eps: 1e-8 lr: 1.0e-5 weight_decay: 0 ``` Refer to the full file for complete details (`!cat llama3_2_1b_squad.yaml`). Finally, [authenticate](https://huggingface.co/docs/hub/en/security-tokens) the VM running the notebook with Hugging Face to download the model and dataset: ```python from getpass import getpass hf_token = getpass("HF token: ") ``` ```bash !hf auth login --token {hf_token} ``` ### Single-Node Because AutoModel is installed through the init script, the `automodel` CLI is available on all nodes. To run training on a single GPU, use this command: ```bash !automodel llama3_2_1b_squad.yaml \ --step_scheduler.max_steps 20 \ --checkpoint.checkpoint_dir /Volumes////checkpoints_single/ \ --checkpoint.is_async True ``` In addition to specifying the configuration file, the command accepts the following options: * `--step_scheduler.max_steps`: Limits the number of training steps taken. This limit is for demonstration purposes. Adapt it for your use case. * `--checkpoint.checkpoint_dir`: Tells AutoModel where to [save model checkpoints](/development/checkpointing) from training. Save the model checkpoints in a Databricks Unity Catalog [volume](https://docs.databricks.com/aws/en/volumes/). * `--checkpoint.staging_dir`: Optional. Specifies a temporary local staging location for checkpoint consolidation before completed files are copied to the final `checkpoint_dir` location. * `--checkpoint.is_async`: Uses asynchronous checkpointing. The GPU metrics in Databricks show that the single GPU has approximately 95% utilization. Single GPU utilization of approximately 95% during model training. To use all four GPUs on this `g6e.12xlarge` instance, add `--nproc-per-node=4` to the command: ```bash !automodel --nproc-per-node=4 llama3_2_1b_squad.yaml \ --step_scheduler.max_steps 20 \ --checkpoint.checkpoint_dir /Volumes////checkpoints_multi/ \ --checkpoint.is_async True ``` The `automodel` CLI uses PyTorch's [Elastic Launch](https://docs.pytorch.org/docs/stable/elastic/run.html) internally to spawn and coordinate multiple training processes on the VM. Each training process runs on a separate GPU, and all four GPUs show approximately 95% utilization. Multi-GPU, single-node utilization of approximately 95% during model training. ### Multi-Node To scale further to multi-node training, submit training jobs to all instances in the Databricks cluster. First, each instance needs to be authenticated with Hugging Face to download the model and dataset: ```python # Ensure workers are authenticated with Hugging Face import subprocess import shlex def run_command(cmd): p = subprocess.run(shlex.split(cmd), capture_output=True) return p.stdout.decode() rdd = sc.parallelize(range(sc.defaultParallelism)) rdd.mapPartitions(lambda _: [run_command("hf auth login --token " + hf_token)]).collect(); ``` Next, use PySpark's [TorchDistributor](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.torch.distributor.TorchDistributor.html) to run the training job across multiple instances: ```py from pyspark.ml.torch.distributor import TorchDistributor import nemo_automodel.recipes.llm.train_ft as recipe_mod num_executor = 2 # Number of workers in cluster num_gpus_per_executor = 4 # Number of GPUs per worker distributor = TorchDistributor( num_processes=num_executor * num_gpus_per_executor, local_mode=False, use_gpu=True, ) train_file = recipe_mod.__file__ args = [ "--config", "llama3_2_1b_squad.yaml", "--step_scheduler.max_steps", "20", "--checkpoint.checkpoint_dir", "/Volumes////checkpoints_dist/", "--checkpoint.is_async", "True", ] distributor.run(train_file, *args) ``` `TorchDistributor` uses `torchrun` internally. Point the distributor at the recipe module directly, rather than the `automodel` CLI, which also wraps `torchrun`. During training, GPU utilization is approximately 95% across all eight GPUs on the worker nodes. ## Track Experiments with MLflow Databricks includes built-in MLflow integration for tracking experiments, logging metrics, and storing artifacts. To use MLflow with AutoModel on Databricks, add the MLflow configuration to your YAML file. ### Configure MLflow Edit your configuration file (e.g., `llama3_2_1b_squad.yaml`) to include the `mlflow` section: ```yaml mlflow: experiment_name: "automodel-databricks-llama3-squad" run_name: "" tracking_uri: "databricks" artifact_location: null tags: platform: "databricks" task: "squad-finetune" model_family: "llama3.2" ``` For Databricks, the key configuration parameters are: * `tracking_uri`: Set to `"databricks"` to use Databricks' managed MLflow tracking server * `experiment_name`: The name of your experiment, which appears in the Databricks workspace * `artifact_location`: Leave as `null` to use default Databricks artifact storage, or specify a Unity Catalog volume path like `/Volumes////mlflow-artifacts` * `tags`: Add custom tags to organize and filter your runs Databricks automatically handles authentication when `tracking_uri` is set to `"databricks"`. No additional credentials are needed. ### Run Training with MLflow Run training with MLflow tracking enabled using the same commands as before. AutoModel reads the MLflow configuration from the YAML file: **Single-node:** ```bash !automodel llama3_2_1b_squad.yaml \ --step_scheduler.max_steps 20 \ --checkpoint.checkpoint_dir /Volumes////checkpoints/ ``` **Multi-GPU:** ```bash !automodel --nproc-per-node=4 llama3_2_1b_squad.yaml \ --step_scheduler.max_steps 20 \ --checkpoint.checkpoint_dir /Volumes////checkpoints/ ``` **Multi-node with TorchDistributor:** ```python import nemo_automodel.recipes.llm.train_ft as recipe_mod distributor = TorchDistributor( num_processes=num_executor * num_gpus_per_executor, local_mode=False, use_gpu=True, ) args = [ "--config", "llama3_2_1b_squad.yaml", "--step_scheduler.max_steps", "20", "--checkpoint.checkpoint_dir", "/Volumes////checkpoints/", ] distributor.run(recipe_mod.__file__, *args) ``` ### View Results During training, you'll see MLflow logging messages in your output: ``` MLflow run started: abc123def456 View run at: databricks/#/mlflow/experiments/123/runs/abc123def456 ``` To view your experiments and metrics: 1. Navigate to the **Experiments** page in your Databricks workspace 2. Find your experiment by name (e.g., `automodel-databricks-llama3-squad`) 3. Click on a run to view metrics, parameters, and artifacts The Databricks MLflow UI displays: * Training and validation metrics over time * Model parameters and hyperparameters * Custom tags for filtering and comparison * Artifacts and model checkpoints * System metrics (GPU utilization, memory usage) ### Store Artifacts in Unity Catalog To store MLflow artifacts in Unity Catalog volumes, specify the `artifact_location`: ```yaml mlflow: experiment_name: "automodel-databricks-llama3-squad" tracking_uri: "databricks" artifact_location: "/Volumes////mlflow-artifacts" tags: platform: "databricks" ``` This ensures your artifacts are stored in a governed, versioned location within Unity Catalog. ### Additional Configuration You can override MLflow settings from the command line: ```bash !automodel llama3_2_1b_squad.yaml \ --mlflow.experiment_name "custom-experiment-name" \ --mlflow.run_name "baseline-run-1" \ --mlflow.tags.learning_rate "1e-5" ``` For more details on MLflow configuration options and best practices, see the [MLflow logging guide](/development/mlflow-logging). ## Conclusion This guide showed how to use AutoModel for model training on Databricks-managed compute, scaling from a single GPU to multi-GPU and multi-node training. Although this guide fine-tunes a Llama-3.2-1B model using the SQuAD dataset, any supported AutoModel functionality, such as model pre-training and vision-language models (VLMs), can run and scale on Databricks. Check [additional recipes and end-to-end examples](/recipes-e2e-examples/overview) to learn more.