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
def create_or_get_secret(name: str, value: str | None, label: str):
6
    if not value:
7
        raise ValueError(f"{label} environment variable is not set. Set it and try again.")
8
    try:
9
        secret = client.secrets.create(
10
            name=name,
11
            workspace="default",
12
            value=value,
13
        )
14
        print(f"Created secret: {name}")
15
        return secret
16
    except ConflictError:
17
        print(f"Secret '{name}' already exists, continuing...")
18
        return client.secrets.retrieve(name=name, workspace="default")
19
20
# Create Hugging Face token secret
21
hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")
22
print("HF_TOKEN secret:")
23
print(hf_secret.model_dump_json(indent=2))
24
25
# Create NGC API key secret
26
# Uncomment the line below if you have NGC API Key and want to finetune NGC models
27
# 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
# Dashboard drawing helpers
33
# ---------------------------------------------------------------------------
34
35
_PALETTE = {
36
    "val_loss":   "#E74C3C",
37
    "train_loss": "#F39C12",
38
    "vram":       ["#3498DB", "#9B59B6", "#1ABC9C", "#E67E22"],
39
    "util":       ["#2ECC71", "#E74C3C", "#3498DB", "#F1C40F"],
40
    "grid":       "#ECECEC",
41
    "title":      "#2C3E50",
42
    "subtitle":   "#7F8C8D",
43
    "spine":      "#CCCCCC",
44
    "tick":       "#666666",
45
}
46
47
def _style_axis(ax):
48
    """Apply shared cosmetic styling to a subplot axis."""
49
    ax.set_facecolor("white")
50
    ax.grid(True, alpha=0.4, color=_PALETTE["grid"], linewidth=0.8)
51
    for spine in ("top", "right"):
52
        ax.spines[spine].set_visible(False)
53
    ax.spines["left"].set_color(_PALETTE["spine"])
54
    ax.spines["bottom"].set_color(_PALETTE["spine"])
55
    ax.tick_params(colors=_PALETTE["tick"], labelsize=9)
56
57
def _plot_line(ax, xs, ys, color, label, fill=True):
58
    """Plot a time series, gracefully skipping None values."""
59
    pts = [(x, y) for x, y in zip(xs, ys) if y is not None]
60
    if not pts:
61
        return
62
    px, py = zip(*pts)
63
    ax.plot(
64
        px, py, color=color, linewidth=2.2,
65
        marker="o", markersize=4,
66
        markerfacecolor="white", markeredgewidth=1.8, markeredgecolor=color,
67
        label=label, zorder=3,
68
    )
69
    if fill:
70
        ax.fill_between(px, py, alpha=0.08, color=color)
71
72
def _plot_gpu_panel(ax, xs, history, colors, fallback_label):
73
    """Plot per-GPU time series with area fill."""
74
    if not history or not history[0]:
75
        ax.text(
76
            0.5, 0.5, "No GPU data", transform=ax.transAxes,
77
            ha="center", va="center", fontsize=11, color="#AAAAAA",
78
        )
79
        return
80
    n_gpus = max(len(snap) for snap in history)
81
    for g in range(n_gpus):
82
        vals = [snap[g] if g < len(snap) else 0 for snap in history]
83
        c = colors[g % len(colors)]
84
        label = f"GPU {g}" if n_gpus > 1 else fallback_label
85
        ax.plot(xs[: len(vals)], vals, color=c, linewidth=2, label=label)
86
        ax.fill_between(xs[: len(vals)], vals, alpha=0.08, color=c)
87
    if n_gpus > 1:
88
        ax.legend(fontsize=9, framealpha=0.9, edgecolor="#DDD")
89
90
def _draw_dashboard(
91
    elapsed_mins, val_losses, train_losses,
92
    vram_history, util_history,
93
    job_name, status_str, step_str, elapsed_str,
94
):
95
    """Render a live 1x3 training dashboard."""
96
    fig, axes = plt.subplots(1, 3, figsize=(20, 5.5))
97
    fig.patch.set_facecolor("#FAFBFC")
98
99
    fig.suptitle(
100
        job_name, fontsize=15, fontweight="bold",
101
        color=_PALETTE["title"], y=1.10,
102
    )
103
    fig.text(
104
        0.5, 1.01,
105
        f"{status_str}   |   {step_str}   |   {elapsed_str}",
106
        ha="center", fontsize=13, color=_PALETTE["subtitle"],
107
    )
108
109
    for ax in axes:
110
        _style_axis(ax)
111
112
    # -- Panel 1: Loss curves --
113
    _plot_line(axes[0], elapsed_mins, val_losses, _PALETTE["val_loss"], "Val Loss", fill=True)
114
    _plot_line(axes[0], elapsed_mins, train_losses, _PALETTE["train_loss"], "Train Loss", fill=False)
115
    axes[0].set_title("Train/Validation Loss", fontsize=13, fontweight="bold", color=_PALETTE["title"], pad=12)
116
    axes[0].set_xlabel("Time (min)", fontsize=10, color="#666")
117
    axes[0].set_ylabel("Loss", fontsize=10, color="#666")
118
    if any(v is not None for v in val_losses + train_losses):
119
        axes[0].legend(fontsize=9, framealpha=0.9, edgecolor="#DDD")
120
121
    # -- Panel 2: GPU VRAM usage --
122
    _plot_gpu_panel(axes[1], elapsed_mins, vram_history, _PALETTE["vram"], "VRAM")
123
    axes[1].set_title("GPU VRAM Usage", fontsize=13, fontweight="bold", color=_PALETTE["title"], pad=12)
124
    axes[1].set_xlabel("Time (min)", fontsize=10, color="#666")
125
    axes[1].set_ylabel("Usage (%)", fontsize=10, color="#666")
126
    axes[1].set_ylim(-2, 105)
127
128
    # -- Panel 3: GPU utilization --
129
    _plot_gpu_panel(axes[2], elapsed_mins, util_history, _PALETTE["util"], "Utilization")
130
    axes[2].set_title("GPU Utilization", fontsize=13, fontweight="bold", color=_PALETTE["title"], pad=12)
131
    axes[2].set_xlabel("Time (min)", fontsize=10, color="#666")
132
    axes[2].set_ylabel("Utilization (%)", fontsize=10, color="#666")
133
    axes[2].set_ylim(-2, 105)
134
135
    plt.tight_layout(rect=[0, 0, 1, 0.98])
136
    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
def get_training_metric(
13
    status: PlatformJobStatusResponse,
14
    metric_key: str,
15
) -> float | None:
16
    """Return a metric reported by a task in the training step."""
17
    for job_step in status.steps or []:
18
        if job_step.name == "training":
19
            for task in job_step.tasks or []:
20
                value = (task.status_details or {}).get(metric_key)
21
                if value is not None:
22
                    return float(value)
23
    return None
24
25
# ---------------------------------------------------------------------------
26
# Job polling with live dashboard
27
# ---------------------------------------------------------------------------
28
29
def wait_for_job(
30
    workspace: str,
31
    job_name: str,
32
    timeout: int = TIMEOUT_SECONDS,
33
    poll_interval: int = 10,
34
    val_loss_key: str = VAL_LOSS_KEY,
35
    train_loss_key: str = TRAIN_LOSS_KEY,
36
) -> PlatformJobStatusResponse:
37
    """
38
    Poll job status until completed, failed, cancelled, or timeout.
39
    Displays a live dashboard with loss curves and GPU metrics.
40
41
    Args:
42
        workspace: The workspace where the job is running.
43
        job_name: The name of the job to monitor.
44
        timeout: Maximum time to wait in seconds (default: 30 minutes).
45
        poll_interval: Time between status checks in seconds (default: 10).
46
47
    Returns:
48
        The final job status response.
49
    """
50
    start_time = time.time()
51
52
    # Time-series accumulators required for plotting
53
    elapsed_mins: list[float] = []
54
    val_losses: list[float | None] = []
55
    train_losses: list[float | None] = []
56
    vram_history: list[list[float]] = []
57
    util_history: list[list[float]] = []
58
59
    while True:
60
        elapsed = time.time() - start_time
61
        elapsed_min = elapsed / 60
62
63
        # Check for timeout
64
        if elapsed > timeout:
65
            error_message = f"Timeout reached after {elapsed_min:.1f} minutes"
66
            print(f"\n{error_message}")
67
            print("Job did not complete within the timeout period.")
68
            raise Exception(error_message)
69
70
        status = client.jobs.get_status(name=job_name, workspace=workspace)
71
72
        # -- Extract training progress from nested steps structure --
73
        step: int | None = None
74
        max_steps: int | None = None
75
        training_phase: str | None = None
76
        val_loss: float | None = None
77
        train_loss: float | None = None
78
        current_step_name: str | None = None
79
        current_step_phase: str | None = None
80
81
        for job_step in status.steps or []:
82
            # Track the current active step name and phase for progress display
83
            if job_step.tasks:
84
                task = job_step.tasks[0]
85
                td = task.status_details or {}
86
                phase = cast(str, td.get("phase", ""))
87
                # Update current step if it's active or pending (not completed)
88
                if job_step.status in ("active", "pending"):
89
                    current_step_name = job_step.name
90
                    current_step_phase = phase or "started"
91
92
            if job_step.name == "training":
93
                for task in job_step.tasks or []:
94
                    td = task.status_details or {}
95
                    step = cast(int, td["step"]) if "step" in td else None
96
                    max_steps = cast(int, td["max_steps"]) if "max_steps" in td else None
97
                    training_phase = cast(str, td["phase"]) if "phase" in td else None
98
                    raw_val_loss = td.get(val_loss_key)
99
                    val_loss = float(raw_val_loss) if raw_val_loss is not None else None
100
                    raw_train_loss = td.get(train_loss_key)
101
                    train_loss = float(raw_train_loss) if raw_train_loss is not None else None
102
                    break
103
                break
104
105
        if val_loss is None:
106
            raw_val_loss = (status.status_details or {}).get(val_loss_key)
107
            val_loss = float(raw_val_loss) if raw_val_loss is not None else None
108
        if train_loss is None:
109
            raw_train_loss = (status.status_details or {}).get(train_loss_key)
110
            train_loss = float(raw_train_loss) if raw_train_loss is not None else None
111
112
        # -- Collect GPU snapshot --
113
        vram_pcts, util_pcts = _get_gpu_snapshot()
114
115
        # -- Append to accumulators used for the plots --
116
        elapsed_mins.append(elapsed_min)
117
        val_losses.append(val_loss)
118
        train_losses.append(train_loss)
119
        vram_history.append(vram_pcts)
120
        util_history.append(util_pcts)
121
122
        # -- Build status strings --
123
        status_str = f"Status: {status.status}"
124
        if step is not None and max_steps is not None:
125
            pct = step / max_steps * 100
126
            step_str = f"Step {step}/{max_steps} ({pct:.0f}%)"
127
            if training_phase:
128
                step_str += f" - {training_phase}"
129
        else:
130
            if current_step_name and current_step_phase:
131
                step_str = f"{current_step_name} - {current_step_phase}"
132
            elif current_step_name:
133
                step_str = f"{current_step_name}"
134
            else:
135
                step_str = "Waiting for training to start..."
136
        elapsed_str = f"Elapsed: {elapsed_min:.1f} min"
137
138
        # -- Redraw dashboard --
139
        clear_output(wait=True)
140
        _draw_dashboard(
141
            elapsed_mins, val_losses, train_losses,
142
            vram_history, util_history,
143
            job_name, status_str, step_str, elapsed_str,
144
        )
145
146
        # -- Check terminal conditions --
147
        if status.status.lower() == "completed":
148
            # Redraw dashboard one final time with "completed" status
149
            status_str = f"Status: {status.status}"
150
            if step is not None and max_steps is not None:
151
                step_str = f"Step {max_steps}/{max_steps} (100%)"
152
            clear_output(wait=True)
153
            _draw_dashboard(
154
                elapsed_mins, val_losses, train_losses,
155
                vram_history, util_history,
156
                job_name, status_str, step_str, elapsed_str,
157
            )
158
            print(f"\nJob completed in {elapsed_min:.1f} minutes ({elapsed:.0f}s)")
159
            return status
160
        elif status.status.lower() in ("failed", "cancelled", "error"):
161
            print(f"\nJob finished with status: {status.status}")
162
            print(f"Total time elapsed: {elapsed_min:.1f} minutes ({elapsed:.0f}s)")
163
164
            # Print error details from the job level
165
            if status.error_details:
166
                error_msg = status.error_details.get("message", "")
167
                if error_msg:
168
                    print(f"\nError: {error_msg}")
169
170
            # Find and print error details from the failed step/task
171
            for job_step in status.steps or []:
172
                if job_step.status == "error":
173
                    print(f"\nFailed step: {job_step.name}")
174
                    if job_step.error_details:
175
                        step_error = job_step.error_details.get("message", "")
176
                        if step_error:
177
                            print(f"Step error: {step_error}")
178
                    # Get error_stack from the failed task
179
                    for task in job_step.tasks or []:
180
                        if task.status == "error" and hasattr(task, "error_stack") and task.error_stack:
181
                            print(f"\nError stack trace:\n{task.error_stack}")
182
                        elif task.status == "error" and task.error_details:
183
                            task_error = task.error_details.get("message", "")
184
                            if task_error:
185
                                print(f"Task error: {task_error}")
186
                    break
187
188
            raise Exception(f"Job finished with status: {status.status}")
189
190
        time.sleep(poll_interval)
191
192
# Wait for the job to complete
193
job_with_sequence_packing_status = wait_for_job(
194
    workspace="default",
195
    job_name=job_with_sequence_packing.job.name,
196
    timeout=TIMEOUT_SECONDS,
197
)
198
199
packed_val_loss = get_training_metric(job_with_sequence_packing_status, VAL_LOSS_KEY)
200
if packed_val_loss is not None:
201
    print(f"Validation loss: {packed_val_loss:.2f}")
202
else:
203
    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')