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

1
from datasets import load_dataset, DatasetDict
2
3
print("Loading dataset rajpurkar/squad")
4
raw_dataset = load_dataset("rajpurkar/squad")
5
if not isinstance(raw_dataset, DatasetDict):
6
    raise ValueError("Dataset does not contain expected splits")
7
8
print("Loaded dataset")
9
10
SEED = 1234
11
TRAINING_SIZE = 3000
12
VALIDATION_SIZE = 300
13
TEST_SIZE = 100
14
DATASET_PATH = Path("kd-dataset").absolute()
15
16
os.makedirs(DATASET_PATH, exist_ok=True)
17
18
train_set = raw_dataset.get('train')
19
split = train_set.train_test_split(test_size=0.05, seed=SEED)
20
21
train_ds = split['train'].select(range(min(TRAINING_SIZE, len(split['train']))))
22
val_ds = split['test'].select(range(min(VALIDATION_SIZE, len(split['test']))))
23
test_ds = split['test'].select(range(VALIDATION_SIZE, min(VALIDATION_SIZE + TEST_SIZE, len(split['test']))))
24
25
26
def convert_squad(example):
27
    """Convert SQuAD format to prompt/completion format."""
28
    prompt = f"Context: {example['context']} Question: {example['question']} Answer:"
29
    completion = example["answers"]["text"][0]
30
    return {"prompt": prompt, "completion": completion}
31
32
33
def write_jsonl(dataset, path):
34
    with open(path, "w", encoding="utf-8") as f:
35
        for example in dataset:
36
            f.write(json.dumps(convert_squad(example)) + "\n")
37
38
39
def write_test_jsonl(dataset, path):
40
    """Save test split with raw context/question for chat-style evaluation."""
41
    with open(path, "w", encoding="utf-8") as f:
42
        for example in dataset:
43
            f.write(json.dumps({
44
                "context": example["context"],
45
                "question": example["question"],
46
                "completion": example["answers"]["text"][0],
47
            }) + "\n")
48
49
50
write_jsonl(train_ds, f"{DATASET_PATH}/training.jsonl")
51
write_jsonl(val_ds, f"{DATASET_PATH}/validation.jsonl")
52
write_test_jsonl(test_ds, f"{DATASET_PATH}/testing.jsonl")
53
54
print(f"Training:   {len(train_ds)} rows")
55
print(f"Validation: {len(val_ds)} rows")
56
print(f"Test:       {len(test_ds)} rows")
57
58
with open(f"{DATASET_PATH}/training.jsonl", 'r') as f:
59
    sample = json.loads(f.readline())
60
    print(f"\nSample prompt:     {sample['prompt'][:150]}...")
61
    print(f"Sample completion:  {sample['completion']}")

1
DATASET_NAME = "kd-dataset"
2
3
try:
4
    client.files.filesets.create(
5
        workspace="default",
6
        name=DATASET_NAME,
7
        description="Knowledge distillation training data"
8
    )
9
    print(f"Created fileset: {DATASET_NAME}")
10
except ConflictError:
11
    print(f"Fileset '{DATASET_NAME}' already exists, continuing...")
12
13
client.files.upload(
14
    local_path=f"{DATASET_PATH}/",
15
    remote_path="",
16
    fileset=DATASET_NAME,
17
    workspace="default"
18
)
19
20
print("Uploaded files:")
21
print(json.dumps([f.model_dump() for f in client.files.list(fileset=DATASET_NAME, workspace="default").data], indent=2))

1
HF_TOKEN = os.getenv("HF_TOKEN")
2
3
4
def create_or_get_secret(name: str, value: str | None, label: str):
5
    if not value:
6
        raise ValueError(f"{label} is not set")
7
    try:
8
        secret = client.secrets.create(
9
            name=name,
10
            workspace="default",
11
            value=value,
12
        )
13
        print(f"Created secret: {name}")
14
        return secret
15
    except ConflictError:
16
        print(f"Secret '{name}' already exists, continuing...")
17
        return client.secrets.retrieve(name=name, workspace="default")
18
19
20
hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")
21
print("HF_TOKEN secret:")
22
print(hf_secret.model_dump_json(indent=2))

1
from nemo_platform.types.files import HuggingfaceStorageConfigParam
2
3
SPEC_TIMEOUT_SECONDS = 120
4
5
6
def create_model(hf_repo: str, model_name: str, description: str):
7
    """Create a fileset + model entity and wait for ModelSpec."""
8
    try:
9
        client.files.filesets.create(
10
            workspace="default",
11
            name=model_name,
12
            description=description,
13
            storage=HuggingfaceStorageConfigParam(
14
                type="huggingface",
15
                repo_id=hf_repo,
16
                repo_type="model",
17
                token_secret=hf_secret.name
18
            )
19
        )
20
        print(f"Created fileset: {model_name}")
21
    except ConflictError:
22
        print(f"Fileset '{model_name}' already exists.")
23
24
    try:
25
        model = client.models.create(
26
            workspace="default",
27
            name=model_name,
28
            fileset=f"default/{model_name}",
29
        )
30
        print(f"Created Model Entity: {model_name}")
31
    except ConflictError:
32
        print(f"Model '{model_name}' already exists. Updating fileset.")
33
        model = client.models.update(
34
            workspace="default",
35
            name=model_name,
36
            fileset=f"default/{model_name}",
37
        )
38
39
    print(f"Waiting for ModelSpec on {model_name}...")
40
    spec_start = time.time()
41
    while not model.spec:
42
        if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:
43
            raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS}s")
44
        time.sleep(2)
45
        model = client.models.retrieve(workspace="default", name=model_name)
46
    print(f"ModelSpec populated: {model.spec}")
47
    return model
48
49
50
student_model = create_model(
51
    hf_repo="meta-llama/Llama-3.2-1B-Instruct",
52
    model_name="llama-3-2-1b-student",
53
    description="Llama 3.2 1B Instruct student model",
54
)
55
56
print()
57
58
teacher_model = create_model(
59
    hf_repo="meta-llama/Llama-3.2-3B-Instruct",
60
    model_name="llama-3-2-3b-teacher",
61
    description="Llama 3.2 3B Instruct teacher model",
62
)

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

1
from IPython.display import clear_output
2
3
4
def wait_for_job(job_name: str):
5
    """Poll job status until completion."""
6
    while True:
7
        status = client.jobs.get_status(name=job_name, workspace="default")
8
        clear_output(wait=True)
9
        print(f"Job: {job_name}")
10
        print(f"Status: {status.status}")
11
12
        for job_step in status.steps or []:
13
            if job_step.name == "training":
14
                for task in job_step.tasks or []:
15
                    details = task.status_details or {}
16
                    step = details.get("step")
17
                    max_steps = details.get("max_steps")
18
                    if step is not None and max_steps is not None:
19
                        print(f"Progress: Step {step}/{max_steps} ({step / max_steps * 100:.1f}%)")
20
                    phase = details.get("phase")
21
                    if phase:
22
                        print(f"Phase: {phase}")
23
                    break
24
                break
25
26
        if status.status in ("completed", "failed", "cancelled", "error"):
27
            print(f"\nJob finished: {status.status}")
28
            return status
29
30
        time.sleep(10)
31
32
33
teacher_status = wait_for_job(TEACHER_JOB_NAME)
34
assert teacher_status.status == "completed"

1
baseline_suffix = uuid.uuid4().hex[:4]
2
BASELINE_DEPLOYMENT_CONFIG = f"baseline-student-cfg-{baseline_suffix}"
3
BASELINE_DEPLOYMENT_NAME = f"baseline-student-{baseline_suffix}"
4
5
baseline_deployment_config = client.inference.deployment_configs.create(
6
    workspace="default",
7
    name=BASELINE_DEPLOYMENT_CONFIG,
8
    engine="vllm",
9
    model_spec={
10
        "model_namespace": "default",
11
        "model_name": student_model.name,
12
    },
13
    executor_config={
14
        "gpu": 1,
15
        "image_name": "vllm/vllm-openai",
16
        "image_tag": "v0.22.1",
17
    },
18
)
19
20
baseline_deployment = client.inference.deployments.create(
21
    workspace="default",
22
    name=BASELINE_DEPLOYMENT_NAME,
23
    config=baseline_deployment_config.name
24
)
25
26
print(f"Baseline student deployment: {baseline_deployment.name}")

1
def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):
2
    """Poll deployment until ready."""
3
    start = time.time()
4
    timeout = timeout_minutes * 60
5
    while True:
6
        dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")
7
        elapsed = time.time() - start
8
        clear_output(wait=True)
9
        print(f"Deployment: {deployment_name}")
10
        print(f"Status: {dep.status}")
11
        print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")
12
13
        if dep.status == "READY":
14
            print("\nDeployment is ready!")
15
            return dep
16
        if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):
17
            print(f"\nDeployment failed: {dep.status}")
18
            return dep
19
        if elapsed > timeout:
20
            print(f"\nTimeout ({timeout_minutes}m). Check status manually.")
21
            return dep
22
        time.sleep(15)
23
24
25
dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)
26
assert dep_status.status == "READY"

1
with open(f"{DATASET_PATH}/testing.jsonl", "r", encoding="utf-8") as f:
2
    test_data = [json.loads(line) for line in f]
3
4
contexts = [row["context"] for row in test_data]
5
questions = [row["question"] for row in test_data]
6
reference_completions = [row["completion"] for row in test_data]
7
8
print(f"Test samples: {len(contexts)}")
9
print(f"Sample context:   {contexts[0]}")
10
print(f"Sample question:  {questions[0]}")
11
print(f"Sample reference: {reference_completions[0]}")

1
def generate_completions(
2
    deployment_name: str,
3
    output_model_name: str,
4
    contexts: list[str],
5
    questions: list[str],
6
) -> list[str]:
7
    """Generate completions for a list of context/question pairs using a deployed model."""
8
    completions = []
9
    for context, question in zip(contexts, questions):
10
        messages = [
11
            {
12
                "role": "user",
13
                "content": f"Based on the following context, answer the question.\n\nContext: {context}\n\nQuestion: {question}",
14
            }
15
        ]
16
        response = client.inference.gateway.provider.post(
17
            "v1/chat/completions",
18
            name=deployment_name,
19
            workspace="default",
20
            body={
21
                "model": f"default/{output_model_name}",
22
                "messages": messages,
23
                "temperature": 0,
24
                "max_tokens": 128,
25
            }
26
        )
27
        completions.append(response["choices"][0]["message"]["content"])
28
    return completions
29
30
31
print("Generating baseline (base student) predictions...")
32
baseline_completions = generate_completions(BASELINE_DEPLOYMENT_NAME, student_model.name, contexts, questions)
33
print(f"Generated {len(baseline_completions)} baseline predictions")
34
print(f"\nSample baseline output: {baseline_completions[0]}")

1
client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace="default")
2
print(f"Deleted baseline deployment: {BASELINE_DEPLOYMENT_NAME}")
3
4
if not client.models.wait_for_status(
5
    deployment_name=BASELINE_DEPLOYMENT_NAME,
6
    desired_status="DELETED",
7
    workspace="default",
8
    timeout=600,
9
):
10
    raise TimeoutError(
11
        f"Deployment {BASELINE_DEPLOYMENT_NAME} was not deleted within timeout"
12
    )
13
14
client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace="default")
15
print(f"Deleted baseline deployment config: {BASELINE_DEPLOYMENT_CONFIG}")

1
from nemo_automodel_plugin.schema import AutomodelJobInput
2
3
KD_JOB_NAME = f"my-kd-job-{job_suffix}"
4
KD_OUTPUT_NAME = f"kd-student-{job_suffix}"
5
6
kd_spec = AutomodelJobInput(
7
    model=f"default/{student_model.name}",
8
    dataset={"training": f"default/{DATASET_NAME}"},
9
    training={
10
        "training_type": "distillation",
11
        "finetuning_type": "all_weights",
12
        "teacher_model": f"default/{TRAINED_TEACHER_NAME}",
13
        "teacher_precision": "bf16",
14
        "distillation_ratio": 0.5,
15
        "distillation_temperature": 2.0,
16
        "max_seq_length": 2048,
17
    },
18
    schedule={"epochs": 1},
19
    batch={"global_batch_size": 64, "micro_batch_size": 1},
20
    optimizer={"learning_rate": 5e-5},
21
    parallelism={"num_gpus_per_node": 1},
22
    output={"name": KD_OUTPUT_NAME},
23
)
24
25
kd_job = client.customization.automodel.jobs.create(
26
    spec=kd_spec, workspace="default", name=KD_JOB_NAME
27
)
28
29
DISTILLED_STUDENT_NAME = KD_OUTPUT_NAME
30
print(f"Distillation job: {kd_job.job.name}")
31
print(f"Output student model: {DISTILLED_STUDENT_NAME}")

1
kd_status = wait_for_job(KD_JOB_NAME)
2
assert kd_status.status == "completed"

1
deploy_suffix_2 = uuid.uuid4().hex[:4]
2
STUDENT_DEPLOYMENT_CONFIG = f"kd-student-deploy-cfg-{deploy_suffix_2}"
3
STUDENT_DEPLOYMENT_NAME = f"kd-student-deploy-{deploy_suffix_2}"
4
5
student_deployment_config = client.inference.deployment_configs.create(
6
    workspace="default",
7
    name=STUDENT_DEPLOYMENT_CONFIG,
8
    engine="vllm",
9
    model_spec={
10
        "model_namespace": "default",
11
        "model_name": DISTILLED_STUDENT_NAME,
12
    },
13
    executor_config={
14
        "gpu": 1,
15
        "image_name": "vllm/vllm-openai",
16
        "image_tag": "v0.22.1",
17
    },
18
)
19
20
student_deployment = client.inference.deployments.create(
21
    workspace="default",
22
    name=STUDENT_DEPLOYMENT_NAME,
23
    config=student_deployment_config.name
24
)
25
26
print(f"Student deployment: {student_deployment.name}")

1
wait_for_deployment(STUDENT_DEPLOYMENT_NAME)

1
print("Generating distilled student predictions...")
2
student_completions = generate_completions(STUDENT_DEPLOYMENT_NAME, DISTILLED_STUDENT_NAME, contexts, questions)
3
print(f"Generated {len(student_completions)} student predictions")
4
print(f"\nSample student output: {student_completions[0]}")

1
import evaluate
2
3
rouge = evaluate.load("rouge")
4
5
baseline_scores = rouge.compute(predictions=baseline_completions, references=reference_completions)
6
student_scores = rouge.compute(predictions=student_completions, references=reference_completions)
7
8
metrics = list(baseline_scores.keys())
9
header = f"{'Model':<35} " + " ".join(f"{m:>10}" for m in metrics)
10
separator = "-" * len(header)
11
12
print("=" * 60)
13
print("ROUGE SCORE COMPARISON")
14
print("=" * 60)
15
print(header)
16
print(separator)
17
print(f"{'Base Student (1B, no training)':<35} " + " ".join(f"{baseline_scores[m]:>10.4f}" for m in metrics))
18
print(f"{'Distilled Student (1B, KD)':<35} " + " ".join(f"{student_scores[m]:>10.4f}" for m in metrics))

1
print("Sample predictions (first 3):\n")
2
for i in range(min(3, len(contexts))):
3
    print(f"--- Sample {i + 1} ---")
4
    print(f"Question:  {questions[i]}")
5
    print(f"Reference: {reference_completions[i]}")
6
    print(f"Baseline:  {baseline_completions[i][:200]}")
7
    print(f"Distilled: {student_completions[i][:200]}")
8
    print()