1
import json
2
import os
3
from nemo_platform import NeMoPlatform, ConflictError
4
5
NMP_BASE_URL = os.environ.get("NMP_BASE_URL", "http://localhost:8080")
6
client = NeMoPlatform(
7
    base_url=NMP_BASE_URL,
8
    workspace="default"
9
)

1
from pathlib import Path
2
from datasets import load_dataset, DatasetDict
3
import json
4
5
# Load the SQuAD dataset from Hugging Face
6
print("Loading dataset rajpurkar/squad")
7
raw_dataset = load_dataset("rajpurkar/squad")
8
if not isinstance(raw_dataset, DatasetDict):
9
    raise ValueError("Dataset does not contain expected splits")
10
11
print("Loaded dataset")
12
13
# Configuration
14
VALIDATION_PROPORTION = 0.05
15
SEED = 1234
16
17
# For the purpose of this tutorial, we'll use a subset of the dataset
18
# The larger the datasets, the better the model will perform but longer the training will take
19
training_size = 3000
20
validation_size = 300
21
DATASET_PATH = Path("sft-dataset").absolute()
22
23
# Create directory if it doesn't exist
24
os.makedirs(DATASET_PATH, exist_ok=True)
25
26
# Get the train split and create a validation split from it
27
train_set = raw_dataset.get('train')
28
split_dataset = train_set.train_test_split(test_size=VALIDATION_PROPORTION, seed=SEED)
29
30
# Select subsets for the tutorial
31
train_ds = split_dataset['train'].select(range(min(training_size, len(split_dataset['train']))))
32
validation_ds = split_dataset['test'].select(range(min(validation_size, len(split_dataset['test']))))
33
34
# Convert SQuAD format to prompt/completion format and save to JSONL
35
def convert_squad_to_sft_format(example):
36
    """Convert SQuAD format to prompt/completion format for SFT training."""
37
    prompt = f"Context: {example['context']} Question: {example['question']} Answer:"
38
    completion = example["answers"]["text"][0]  # Take the first answer
39
    return {"prompt": prompt, "completion": completion}
40
41
# Save training data
42
with open(f"{DATASET_PATH}/training.jsonl", "w", encoding="utf-8") as f:
43
    for example in train_ds:
44
        converted = convert_squad_to_sft_format(example)
45
        f.write(json.dumps(converted) + "\n")
46
47
# Save validation data
48
with open(f"{DATASET_PATH}/validation.jsonl", "w", encoding="utf-8") as f:
49
    for example in validation_ds:
50
        converted = convert_squad_to_sft_format(example)
51
        f.write(json.dumps(converted) + "\n")
52
53
print(f"Saved training.jsonl with {len(train_ds)} rows")
54
print(f"Saved validation.jsonl with {len(validation_ds)} rows")
55
56
# Show a sample from the training data
57
print("\nSample from training data:")
58
with open(f"{DATASET_PATH}/training.jsonl", 'r') as f:
59
    first_line = f.readline()
60
    sample = json.loads(first_line)
61
    print(f"Prompt: {sample['prompt'][:200]}...")
62
    print(f"Completion: {sample['completion']}")

1
# Create fileset to store SFT training data
2
DATASET_NAME = "sft-dataset"
3
4
try:
5
    client.files.filesets.create(
6
        workspace="default",
7
        name=DATASET_NAME,
8
        description="SFT training data"
9
    )
10
    print(f"Created fileset: {DATASET_NAME}")
11
except ConflictError:
12
    print(f"Fileset '{DATASET_NAME}' already exists, continuing...")
13
14
# Upload training data files individually to ensure correct structure
15
client.files.upload(
16
    local_path=f"{DATASET_PATH}/",  # Trailing slash uploads directory contents to fileset root
17
    remote_path="",
18
    fileset=DATASET_NAME,
19
    workspace="default"
20
)
21
22
# Validate training data is uploaded correctly
23
print("Training data:")
24
print(json.dumps([f.model_dump() for f in client.files.list(fileset=DATASET_NAME, workspace="default").data], indent=2))

1
# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set
2
HF_TOKEN = os.getenv("HF_TOKEN")
3
NGC_API_KEY = os.getenv("NGC_API_KEY")
4
5
6
def create_or_get_secret(name: str, value: str | None, label: str):
7
    if not value:
8
        raise ValueError(f"{label} is not set")
9
    try:
10
        secret = client.secrets.create(
11
            name=name,
12
            workspace="default",
13
            value=value,
14
        )
15
        print(f"Created secret: {name}")
16
        return secret
17
    except ConflictError:
18
        print(f"Secret '{name}' already exists, continuing...")
19
        return client.secrets.retrieve(name=name, workspace="default")
20
21
22
# Create HuggingFace token secret
23
hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")
24
print("HF_TOKEN secret:")
25
print(hf_secret.model_dump_json(indent=2))
26
27
# Create NGC API key secret
28
# Uncomment the line below if you have NGC API Key and want to finetune NGC models
29
# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")

1
import time
2
3
# Create a fileset pointing to the desired HuggingFace model
4
from nemo_platform.types.files import HuggingfaceStorageConfigParam
5
6
HF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"
7
MODEL_NAME = "llama-3-2-1b-base"
8
9
# Ensure you have a HuggingFace token secret created
10
try:
11
    base_model_fs = client.files.filesets.create(
12
        workspace="default",
13
        name=MODEL_NAME,
14
        description="Llama 3.2 1B base model from HuggingFace",
15
        storage=HuggingfaceStorageConfigParam(
16
            type="huggingface",
17
            # repo_id is the full model name from Hugging Face
18
            repo_id=HF_REPO_ID,
19
            repo_type="model",
20
            # we use the secret created in the previous step
21
            token_secret=hf_secret.name
22
        )
23
    )
24
    print(f"Created base model fileset: {MODEL_NAME}")
25
except ConflictError:
26
    print(f"Base model fileset already exists. Skipping creation.")
27
    base_model_fs = client.files.filesets.retrieve(
28
        workspace="default",
29
        name=MODEL_NAME,
30
    )
31
32
# Create the Model Entity representation.
33
try:
34
    base_model = client.models.create(
35
        workspace="default",
36
        name=MODEL_NAME,
37
        fileset=f"default/{MODEL_NAME}",
38
    )
39
    print(f"Created Model Entity: {MODEL_NAME}")
40
except ConflictError:
41
    print(f"Base model already exists. Updating fileset if different.")
42
    base_model = client.models.update(
43
        workspace="default",
44
        name=MODEL_NAME,
45
        fileset=f"default/{MODEL_NAME}",
46
    )
47
48
print(f"\nBase model fileset: fileset://default/{base_model.name}")
49
print("Base model fileset files list:")
50
print(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))
51
52
# Wait for ModelSpec to be populated from the checkpoint
53
print("\nWaiting for ModelSpec to be populated...")
54
SPEC_TIMEOUT_SECONDS = 120
55
spec_start = time.time()
56
while not base_model.spec:
57
    if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:
58
        raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")
59
    time.sleep(2)
60
    base_model = client.models.retrieve(
61
        workspace="default",
62
        name=MODEL_NAME,
63
    )
64
65
print(f"ModelSpec populated: {base_model.spec}")

1
import uuid
2
from nemo_automodel_plugin.schema import AutomodelJobInput
3
4
job_suffix = uuid.uuid4().hex[:4]
5
6
JOB_NAME = f"my-sft-job-{job_suffix}"
7
OUTPUT_NAME = f"sft-model-{job_suffix}"
8
9
spec = AutomodelJobInput(
10
    model=f"default/{base_model.name}",
11
    dataset={"training": f"default/{DATASET_NAME}"},
12
    training={
13
        "training_type": "sft",
14
        "finetuning_type": "all_weights",
15
        "max_seq_length": 2048,
16
    },
17
    schedule={"epochs": 2},
18
    batch={"global_batch_size": 64, "micro_batch_size": 1},
19
    optimizer={"learning_rate": 5e-5},
20
    parallelism={
21
        "num_gpus_per_node": 1,
22
        "num_nodes": 1,
23
        "tensor_parallel_size": 1,
24
        "pipeline_parallel_size": 1,
25
    },
26
    output={"name": OUTPUT_NAME},
27
)
28
29
job = client.customization.automodel.jobs.create(
30
    spec=spec, workspace="default", name=JOB_NAME
31
)
32
33
print(f"Submitted job: {job.job.name}")
34
print(f"Output model: {OUTPUT_NAME}")

1
import time
2
from IPython.display import clear_output
3
4
# Poll job status every 10 seconds until completed
5
while True:
6
    status = client.jobs.get_status(
7
        name=job.job.name,
8
        workspace="default"
9
    )
10
11
    clear_output(wait=True)
12
    print(f"Job Status: {status.model_dump_json(indent=2)}")
13
14
    # Extract training progress from nested steps structure
15
    step: int | None = None
16
    max_steps: int | None = None
17
    training_phase: str | None = None
18
19
    for job_step in status.steps or []:
20
        if job_step.name == "training":
21
            for task in job_step.tasks or []:
22
                task_details = task.status_details or {}
23
                step = task_details.get("step")
24
                max_steps = task_details.get("max_steps")
25
                training_phase = task_details.get("phase")
26
                break
27
            break
28
29
    if step is not None and max_steps is not None:
30
        progress_pct = (step / max_steps) * 100
31
        print(f"Training Progress: Step {step}/{max_steps} ({progress_pct:.1f}%)")
32
        if training_phase:
33
            print(f"Training Phase: {training_phase}")
34
    else:
35
        print("Training step not started yet or progress info not available")
36
37
    # Exit loop when job reaches a terminal status
38
    if status.status in ("completed", "failed", "cancelled", "error"):
39
        print(f"\nJob finished with status: {status.status}")
40
        break
41
42
    time.sleep(10)
43
44
if status.status != "completed":
45
    raise RuntimeError(f"Training job finished with status: {status.status}")

1
# Validate model entity exists
2
model_entity = client.models.retrieve(workspace='default', name=OUTPUT_NAME)
3
print(model_entity.model_dump_json(indent=2))

1
# Create deployment config
2
deploy_suffix = uuid.uuid4().hex[:4]
3
DEPLOYMENT_CONFIG_NAME = f"sft-model-deployment-cfg-{deploy_suffix}"
4
DEPLOYMENT_NAME = f"sft-model-deployment-{deploy_suffix}"
5
6
deployment_config = client.inference.deployment_configs.create(
7
    workspace="default",
8
    name=DEPLOYMENT_CONFIG_NAME,
9
    engine="vllm",
10
    model_spec={
11
        "model_namespace": "default",
12
        "model_name": OUTPUT_NAME,
13
    },
14
    executor_config={
15
        "gpu": 1,
16
        "image_name": "vllm/vllm-openai",
17
        "image_tag": "v0.22.1",
18
    },
19
)
20
21
# Deploy model using deployment_config created above
22
deployment = client.inference.deployments.create(
23
    workspace="default",
24
    name=DEPLOYMENT_NAME,
25
    config=deployment_config.name
26
)
27
28
29
# Check deployment status
30
deployment_status = client.inference.deployments.retrieve(
31
    name=deployment.name,
32
    workspace="default"
33
)
34
35
print(f"Deployment name: {deployment.name}")
36
print(f"Deployment status: {deployment_status.status}")

1
deployment_config = client.inference.deployment_configs.create(
2
    workspace="default",
3
    name="sft-model-config-multigpu",
4
    engine="vllm",
5
    model_spec={
6
        "model_namespace": "default",
7
        "model_name": OUTPUT_NAME,
8
    },
9
    executor_config={
10
        "gpu": 2,
11
        "image_name": "vllm/vllm-openai",
12
        "image_tag": "v0.22.1",
13
    },
14
)

1
import time
2
from IPython.display import clear_output
3
4
# Poll deployment status every 15 seconds until ready
5
TIMEOUT_MINUTES = 30
6
start_time = time.time()
7
timeout_seconds = TIMEOUT_MINUTES * 60
8
9
print(f"Monitoring deployment '{deployment.name}'...")
10
print(f"Timeout: {TIMEOUT_MINUTES} minutes\n")
11
12
while True:
13
    deployment_status = client.inference.deployments.retrieve(
14
        name=deployment.name,
15
        workspace="default"
16
    )
17
18
    elapsed = time.time() - start_time
19
    elapsed_min = int(elapsed // 60)
20
    elapsed_sec = int(elapsed % 60)
21
22
    clear_output(wait=True)
23
    print(f"Deployment: {deployment.name}")
24
    print(f"Status: {deployment_status.status}")
25
    print(f"Elapsed time: {elapsed_min}m {elapsed_sec}s")
26
27
    # Check if deployment is ready
28
    if deployment_status.status == "READY":
29
        print("\nDeployment is ready!")
30
        if not client.models.wait_for_gateway(deployment.name, workspace="default", timeout=60):
31
            raise RuntimeError("Inference gateway did not become ready")
32
        break
33
34
    # Check for failure states
35
    if deployment_status.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):
36
        raise RuntimeError(f"Deployment failed with status: {deployment_status.status}")
37
38
    # Check timeout
39
    if elapsed > timeout_seconds:
40
        raise TimeoutError(f"Deployment timeout after {TIMEOUT_MINUTES} minutes")
41
42
    time.sleep(15)

1
# Wait for deployment to be ready, then test
2
# Test the fine-tuned model with a question answering prompt
3
context = "The Apollo 11 mission was the first manned mission to land on the Moon. It was launched on July 16, 1969, and Neil Armstrong became the first person to walk on the lunar surface on July 20, 1969. Buzz Aldrin joined him shortly after, while Michael Collins remained in lunar orbit."
4
question = "Who was the first person to walk on the Moon?"
5
6
prompt = f"Context: {context} Question: {question} Answer:"
7
8
response = client.inference.gateway.provider.post(
9
    "v1/completions",
10
    name=deployment.name,
11
    workspace="default",
12
    body={
13
        "model": f"default/{OUTPUT_NAME}",
14
        "prompt": prompt,
15
        "temperature": 0,
16
        "max_tokens": 128
17
    }
18
)
19
20
print("=" * 60)
21
print("MODEL EVALUATION")
22
print("=" * 60)
23
print(f"Question: {question}")
24
print(f"Expected: Neil Armstrong")
25
print(f"Model output: {response['choices'][0]['text']}")