Optimize for Tokens/GPU Throughput

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
import json
2
import os
3
from pathlib import Path
4
from datasets import load_dataset, Dataset, DatasetDict
5
6
# Configuration
7
SEED = 1234
8
DATASET_NAME = "sft-dataset"
9
10
# Convert SQuAD format to prompt/completion format and save to JSONL
11
def convert_squad_to_sft_format(example):
12
    """Convert SQuAD format to prompt/completion format for SFT training."""
13
    prompt = f"Context: {example['context']} Question: {example['question']} Answer:"
14
    completion = example["answers"]["text"][0]  # Take the first answer
15
    return {"prompt": prompt, "completion": completion}
16
17
# Load the SQuAD dataset from Hugging Face
18
print("Loading dataset rajpurkar/squad")
19
ds = load_dataset("rajpurkar/squad")
20
if not isinstance(ds, DatasetDict):
21
    raise ValueError("Dataset does not contain expected splits")
22
23
print("Loaded dataset")
24
25
# For the purpose of this tutorial, we'll use a subset of the dataset
26
# We use a reduced dataset size (3000 training/300 validation samples) to keep tutorial runtime manageable
27
# while still demonstrating the performance benefits of sequence packing. The larger the dataset,
28
# the better the model will perform but the longer the training will take.
29
training_size = 3000
30
validation_size = 300
31
DATASET_PATH = Path(DATASET_NAME).absolute()
32
33
# Get training split and verify it's a Dataset (not IterableDataset)
34
train_dataset = ds["train"]
35
validation_dataset = ds["validation"]
36
assert isinstance(train_dataset, Dataset), "Expected Dataset type"
37
assert isinstance(validation_dataset, Dataset), "Expected Dataset type"
38
39
# Select subsets and save to JSONL files
40
training_ds = train_dataset.select(range(training_size))
41
validation_ds = validation_dataset.select(range(validation_size))
42
43
# Transform to SFT format (prompt/completion)
44
training_ds = training_ds.map(convert_squad_to_sft_format, remove_columns=training_ds.column_names)
45
validation_ds = validation_ds.map(convert_squad_to_sft_format, remove_columns=validation_ds.column_names)
46
47
# Create directory if it doesn't exist
48
# Note: This will create a local 'sft-dataset/' directory with training.jsonl and validation.jsonl files
49
os.makedirs(DATASET_PATH, exist_ok=True)
50
51
# Save subsets to JSONL files
52
training_ds.to_json(f"{DATASET_PATH}/training.jsonl")
53
validation_ds.to_json(f"{DATASET_PATH}/validation.jsonl")
54
55
print(f"Saved training.jsonl with {len(training_ds)} rows")
56
print(f"Saved validation.jsonl with {len(validation_ds)} rows")

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

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

1
# Helpers to draw GPU VRAM Utilization and Validation Loss
2
import matplotlib.pyplot as plt
3
try:
4
    import pynvml
5
    _PYNVML_AVAILABLE = True
6
except ImportError:
7
    _PYNVML_AVAILABLE = False
8
    print("Note: Install nvidia-ml-py ('pip install nvidia-ml-py' or 'uv pip install nvidia-ml-py') to enable live GPU metrics.")
9
10
# ---------------------------------------------------------------------------
11
# GPU metrics collection (nvidia-ml-py; import name is pynvml)
12
# ---------------------------------------------------------------------------
13
14
def _get_gpu_snapshot() -> tuple[list[float], list[float]]:
15
    """Return (vram_usage_pcts, compute_util_pcts) for each GPU."""
16
    if not _PYNVML_AVAILABLE:
17
        return [], []
18
    pynvml.nvmlInit()
19
    try:
20
        vram, util = [], []
21
        for i in range(pynvml.nvmlDeviceGetCount()):
22
            h = pynvml.nvmlDeviceGetHandleByIndex(i)
23
            mem = pynvml.nvmlDeviceGetMemoryInfo(h)
24
            rates = pynvml.nvmlDeviceGetUtilizationRates(h)
25
            vram.append(int(mem.used) / int(mem.total) * 100)
26
            util.append(float(rates.gpu))
27
        return vram, util
28
    finally:
29
        pynvml.nvmlShutdown()
30
31
32
# ---------------------------------------------------------------------------
33
# Dashboard drawing helpers
34
# ---------------------------------------------------------------------------
35
36
_PALETTE = {
37
    "val_loss":   "#E74C3C",
38
    "train_loss": "#F39C12",
39
    "vram":       ["#3498DB", "#9B59B6", "#1ABC9C", "#E67E22"],
40
    "util":       ["#2ECC71", "#E74C3C", "#3498DB", "#F1C40F"],
41
    "grid":       "#ECECEC",
42
    "title":      "#2C3E50",
43
    "subtitle":   "#7F8C8D",
44
    "spine":      "#CCCCCC",
45
    "tick":       "#666666",
46
}
47
48
49
def _style_axis(ax):
50
    """Apply shared cosmetic styling to a subplot axis."""
51
    ax.set_facecolor("white")
52
    ax.grid(True, alpha=0.4, color=_PALETTE["grid"], linewidth=0.8)
53
    for spine in ("top", "right"):
54
        ax.spines[spine].set_visible(False)
55
    ax.spines["left"].set_color(_PALETTE["spine"])
56
    ax.spines["bottom"].set_color(_PALETTE["spine"])
57
    ax.tick_params(colors=_PALETTE["tick"], labelsize=9)
58
59
60
def _plot_line(ax, xs, ys, color, label, fill=True):
61
    """Plot a time series, gracefully skipping None values."""
62
    pts = [(x, y) for x, y in zip(xs, ys) if y is not None]
63
    if not pts:
64
        return
65
    px, py = zip(*pts)
66
    ax.plot(
67
        px, py, color=color, linewidth=2.2,
68
        marker="o", markersize=4,
69
        markerfacecolor="white", markeredgewidth=1.8, markeredgecolor=color,
70
        label=label, zorder=3,
71
    )
72
    if fill:
73
        ax.fill_between(px, py, alpha=0.08, color=color)
74
75
76
def _plot_gpu_panel(ax, xs, history, colors, fallback_label):
77
    """Plot per-GPU time series with area fill."""
78
    if not history or not history[0]:
79
        ax.text(
80
            0.5, 0.5, "No GPU data", transform=ax.transAxes,
81
            ha="center", va="center", fontsize=11, color="#AAAAAA",
82
        )
83
        return
84
    n_gpus = max(len(snap) for snap in history)
85
    for g in range(n_gpus):
86
        vals = [snap[g] if g < len(snap) else 0 for snap in history]
87
        c = colors[g % len(colors)]
88
        label = f"GPU {g}" if n_gpus > 1 else fallback_label
89
        ax.plot(xs[: len(vals)], vals, color=c, linewidth=2, label=label)
90
        ax.fill_between(xs[: len(vals)], vals, alpha=0.08, color=c)
91
    if n_gpus > 1:
92
        ax.legend(fontsize=9, framealpha=0.9, edgecolor="#DDD")
93
94
95
def _draw_dashboard(
96
    elapsed_mins, val_losses, train_losses,
97
    vram_history, util_history,
98
    job_name, status_str, step_str, elapsed_str,
99
):
100
    """Render a live 1x3 training dashboard."""
101
    fig, axes = plt.subplots(1, 3, figsize=(20, 5.5))
102
    fig.patch.set_facecolor("#FAFBFC")
103
104
    fig.suptitle(
105
        job_name, fontsize=15, fontweight="bold",
106
        color=_PALETTE["title"], y=1.10,
107
    )
108
    fig.text(
109
        0.5, 1.01,
110
        f"{status_str}   |   {step_str}   |   {elapsed_str}",
111
        ha="center", fontsize=13, color=_PALETTE["subtitle"],
112
    )
113
114
    for ax in axes:
115
        _style_axis(ax)
116
117
    # -- Panel 1: Loss curves --
118
    _plot_line(axes[0], elapsed_mins, val_losses, _PALETTE["val_loss"], "Val Loss", fill=True)
119
    _plot_line(axes[0], elapsed_mins, train_losses, _PALETTE["train_loss"], "Train Loss", fill=False)
120
    axes[0].set_title("Train/Validation Loss", fontsize=13, fontweight="bold", color=_PALETTE["title"], pad=12)
121
    axes[0].set_xlabel("Time (min)", fontsize=10, color="#666")
122
    axes[0].set_ylabel("Loss", fontsize=10, color="#666")
123
    if any(v is not None for v in val_losses + train_losses):
124
        axes[0].legend(fontsize=9, framealpha=0.9, edgecolor="#DDD")
125
126
    # -- Panel 2: GPU VRAM usage --
127
    _plot_gpu_panel(axes[1], elapsed_mins, vram_history, _PALETTE["vram"], "VRAM")
128
    axes[1].set_title("GPU VRAM Usage", fontsize=13, fontweight="bold", color=_PALETTE["title"], pad=12)
129
    axes[1].set_xlabel("Time (min)", fontsize=10, color="#666")
130
    axes[1].set_ylabel("Usage (%)", fontsize=10, color="#666")
131
    axes[1].set_ylim(-2, 105)
132
133
    # -- Panel 3: GPU utilization --
134
    _plot_gpu_panel(axes[2], elapsed_mins, util_history, _PALETTE["util"], "Utilization")
135
    axes[2].set_title("GPU Utilization", fontsize=13, fontweight="bold", color=_PALETTE["title"], pad=12)
136
    axes[2].set_xlabel("Time (min)", fontsize=10, color="#666")
137
    axes[2].set_ylabel("Utilization (%)", fontsize=10, color="#666")
138
    axes[2].set_ylim(-2, 105)
139
140
    plt.tight_layout(rect=[0, 0, 1, 0.98])
141
    plt.show()

1
import time
2
from typing import cast
3
from IPython.display import clear_output
4
from nemo_platform.types.shared import PlatformJobStatusResponse
5
6
# Timeout set to 30 minutes to accommodate typical LoRA training duration for this dataset size.
7
# Actual training time will vary based on hardware, model size, and dataset complexity.
8
TIMEOUT_SECONDS = 30 * 60  # 30 minutes
9
VAL_LOSS_KEY = "val_loss"
10
TRAIN_LOSS_KEY = "train_loss"
11
12
13
def get_training_metric(
14
    status: PlatformJobStatusResponse,
15
    metric_key: str,
16
) -> float | None:
17
    """Return a metric reported by a task in the training step."""
18
    for job_step in status.steps or []:
19
        if job_step.name == "training":
20
            for task in job_step.tasks or []:
21
                value = (task.status_details or {}).get(metric_key)
22
                if value is not None:
23
                    return float(value)
24
    return None
25
26
27
# ---------------------------------------------------------------------------
28
# Job polling with live dashboard
29
# ---------------------------------------------------------------------------
30
31
def wait_for_job(
32
    workspace: str,
33
    job_name: str,
34
    timeout: int = TIMEOUT_SECONDS,
35
    poll_interval: int = 10,
36
    val_loss_key: str = VAL_LOSS_KEY,
37
    train_loss_key: str = TRAIN_LOSS_KEY,
38
) -> PlatformJobStatusResponse:
39
    """
40
    Poll job status until completed, failed, cancelled, or timeout.
41
    Displays a live dashboard with loss curves and GPU metrics.
42
43
    Args:
44
        workspace: The workspace where the job is running.
45
        job_name: The name of the job to monitor.
46
        timeout: Maximum time to wait in seconds (default: 30 minutes).
47
        poll_interval: Time between status checks in seconds (default: 10).
48
49
    Returns:
50
        The final job status response.
51
    """
52
    start_time = time.time()
53
54
    # Time-series accumulators required for plotting
55
    elapsed_mins: list[float] = []
56
    val_losses: list[float | None] = []
57
    train_losses: list[float | None] = []
58
    vram_history: list[list[float]] = []
59
    util_history: list[list[float]] = []
60
61
    while True:
62
        elapsed = time.time() - start_time
63
        elapsed_min = elapsed / 60
64
65
        # Check for timeout
66
        if elapsed > timeout:
67
            error_message = f"Timeout reached after {elapsed_min:.1f} minutes"
68
            print(f"\n{error_message}")
69
            print("Job did not complete within the timeout period.")
70
            raise Exception(error_message)
71
72
        status = client.jobs.get_status(name=job_name, workspace=workspace)
73
74
        # -- Extract training progress from nested steps structure --
75
        step: int | None = None
76
        max_steps: int | None = None
77
        training_phase: str | None = None
78
        val_loss: float | None = None
79
        train_loss: float | None = None
80
        current_step_name: str | None = None
81
        current_step_phase: str | None = None
82
83
        for job_step in status.steps or []:
84
            # Track the current active step name and phase for progress display
85
            if job_step.tasks:
86
                task = job_step.tasks[0]
87
                td = task.status_details or {}
88
                phase = cast(str, td.get("phase", ""))
89
                # Update current step if it's active or pending (not completed)
90
                if job_step.status in ("active", "pending"):
91
                    current_step_name = job_step.name
92
                    current_step_phase = phase or "started"
93
94
            if job_step.name == "training":
95
                for task in job_step.tasks or []:
96
                    td = task.status_details or {}
97
                    step = cast(int, td["step"]) if "step" in td else None
98
                    max_steps = cast(int, td["max_steps"]) if "max_steps" in td else None
99
                    training_phase = cast(str, td["phase"]) if "phase" in td else None
100
                    raw_val_loss = td.get(val_loss_key)
101
                    val_loss = float(raw_val_loss) if raw_val_loss is not None else None
102
                    raw_train_loss = td.get(train_loss_key)
103
                    train_loss = float(raw_train_loss) if raw_train_loss is not None else None
104
                    break
105
                break
106
107
        if val_loss is None:
108
            raw_val_loss = (status.status_details or {}).get(val_loss_key)
109
            val_loss = float(raw_val_loss) if raw_val_loss is not None else None
110
        if train_loss is None:
111
            raw_train_loss = (status.status_details or {}).get(train_loss_key)
112
            train_loss = float(raw_train_loss) if raw_train_loss is not None else None
113
114
        # -- Collect GPU snapshot --
115
        vram_pcts, util_pcts = _get_gpu_snapshot()
116
117
        # -- Append to accumulators used for the plots --
118
        elapsed_mins.append(elapsed_min)
119
        val_losses.append(val_loss)
120
        train_losses.append(train_loss)
121
        vram_history.append(vram_pcts)
122
        util_history.append(util_pcts)
123
124
        # -- Build status strings --
125
        status_str = f"Status: {status.status}"
126
        if step is not None and max_steps is not None:
127
            pct = step / max_steps * 100
128
            step_str = f"Step {step}/{max_steps} ({pct:.0f}%)"
129
            if training_phase:
130
                step_str += f" - {training_phase}"
131
        else:
132
            if current_step_name and current_step_phase:
133
                step_str = f"{current_step_name} - {current_step_phase}"
134
            elif current_step_name:
135
                step_str = f"{current_step_name}"
136
            else:
137
                step_str = "Waiting for training to start..."
138
        elapsed_str = f"Elapsed: {elapsed_min:.1f} min"
139
140
        # -- Redraw dashboard --
141
        clear_output(wait=True)
142
        _draw_dashboard(
143
            elapsed_mins, val_losses, train_losses,
144
            vram_history, util_history,
145
            job_name, status_str, step_str, elapsed_str,
146
        )
147
148
        # -- Check terminal conditions --
149
        if status.status.lower() == "completed":
150
            # Redraw dashboard one final time with "completed" status
151
            status_str = f"Status: {status.status}"
152
            if step is not None and max_steps is not None:
153
                step_str = f"Step {max_steps}/{max_steps} (100%)"
154
            clear_output(wait=True)
155
            _draw_dashboard(
156
                elapsed_mins, val_losses, train_losses,
157
                vram_history, util_history,
158
                job_name, status_str, step_str, elapsed_str,
159
            )
160
            print(f"\nJob completed in {elapsed_min:.1f} minutes ({elapsed:.0f}s)")
161
            return status
162
        elif status.status.lower() in ("failed", "cancelled", "error"):
163
            print(f"\nJob finished with status: {status.status}")
164
            print(f"Total time elapsed: {elapsed_min:.1f} minutes ({elapsed:.0f}s)")
165
166
            # Print error details from the job level
167
            if status.error_details:
168
                error_msg = status.error_details.get("message", "")
169
                if error_msg:
170
                    print(f"\nError: {error_msg}")
171
172
            # Find and print error details from the failed step/task
173
            for job_step in status.steps or []:
174
                if job_step.status == "error":
175
                    print(f"\nFailed step: {job_step.name}")
176
                    if job_step.error_details:
177
                        step_error = job_step.error_details.get("message", "")
178
                        if step_error:
179
                            print(f"Step error: {step_error}")
180
                    # Get error_stack from the failed task
181
                    for task in job_step.tasks or []:
182
                        if task.status == "error" and hasattr(task, "error_stack") and task.error_stack:
183
                            print(f"\nError stack trace:\n{task.error_stack}")
184
                        elif task.status == "error" and task.error_details:
185
                            task_error = task.error_details.get("message", "")
186
                            if task_error:
187
                                print(f"Task error: {task_error}")
188
                    break
189
190
            raise Exception(f"Job finished with status: {status.status}")
191
192
        time.sleep(poll_interval)
193
194
195
# Wait for the job to complete
196
job_with_sequence_packing_status = wait_for_job(
197
    workspace="default",
198
    job_name=job_with_sequence_packing.job.name,
199
    timeout=TIMEOUT_SECONDS,
200
)
201
202
packed_val_loss = get_training_metric(job_with_sequence_packing_status, VAL_LOSS_KEY)
203
if packed_val_loss is not None:
204
    print(f"Validation loss: {packed_val_loss:.2f}")
205
else:
206
    print("Validation loss: not reported in job status")

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

1
# Wait for the training step to complete
2
job_without_sequence_packing_status = wait_for_job(
3
    workspace="default",
4
    job_name=job_without_sequence_packing.job.name,
5
    timeout=TIMEOUT_SECONDS
6
)
7
8
no_pack_val_loss = get_training_metric(job_without_sequence_packing_status, VAL_LOSS_KEY)
9
if no_pack_val_loss is not None:
10
    print(f"Validation loss: {no_pack_val_loss:.2f}")
11
else:
12
    print("Validation loss: not reported in job status")

1
from nemo_platform.types.jobs import PlatformJobStep
2
import pandas as pd
3
4
STEP_NAME = "training"
5
6
def get_elapsed_time(step: PlatformJobStep) -> float:
7
    """Calculate elapsed time in seconds from step's created_at to updated_at."""
8
    if step.created_at is None or step.updated_at is None:
9
        raise ValueError("Training step timestamps are unavailable")
10
    return (step.updated_at - step.created_at).total_seconds()
11
12
step_with_sequence_packing = client.jobs.steps.retrieve(
13
            name=STEP_NAME,
14
            workspace="default",
15
            job=job_with_sequence_packing.job.name,
16
)
17
18
step_without_sequence_packing = client.jobs.steps.retrieve(
19
            name=STEP_NAME,
20
            workspace="default",
21
            job=job_without_sequence_packing.job.name,
22
)
23
24
time_to_complete_with_sequence_packing = get_elapsed_time(step_with_sequence_packing)
25
time_to_complete_without_sequence_packing = get_elapsed_time(step_without_sequence_packing)
26
27
# Display results as a table
28
results_df = pd.DataFrame({
29
    "Seq Packing Enabled": [True, False],
30
    "Val Loss": [packed_val_loss, no_pack_val_loss],
31
    "Training Step Time, sec": [
32
        time_to_complete_with_sequence_packing,
33
        time_to_complete_without_sequence_packing
34
    ]
35
})
36
37
results_df.style.format({"Val Loss": "{:.2f}", "Training Step Time, sec": "{:.0f}"}).hide(axis='index')