Your GPUs Are Idle: Profiling SageMaker AI Training Jobs and Fixing the Bottleneck

Almost every "we need more GPUs" conversation we walk into on a client engagement ends the same way: the cluster is not short of GPUs, it is short of busy GPUs. A p5 or g6e training job that sits at 35% SM utilization is paying full price for a third of a machine, and nothing in the SageMaker AI console shouts about it. This tutorial shows how to find out what your training job is actually doing, in three escalating levels of effort, and how to fix the three problems that cause most of the idle time.

Everything below runs on SageMaker AI training jobs (the Estimator / PyTorch classes in the SageMaker Python SDK v2). No HyperPod cluster required, though the same techniques apply there.

Level 0: the metrics you already have

Before you instrument anything, read the metrics SageMaker AI publishes for free. Every training job emits GPUUtilization, GPUMemoryUtilization, CPUUtilization, and DiskUtilization to the /aws/sagemaker/TrainingJobs CloudWatch namespace, at one-minute resolution.

import boto3, datetime as dt

cw = boto3.client("cloudwatch")
job = "fsdp-llama-8b-2026-03-04-11-20-31-402"

resp = cw.get_metric_statistics(
    Namespace="/aws/sagemaker/TrainingJobs",
    MetricName="GPUUtilization",
    Dimensions=[{"Name": "Host", "Value": f"{job}/algo-1"}],
    StartTime=dt.datetime.utcnow() - dt.timedelta(hours=6),
    EndTime=dt.datetime.utcnow(),
    Period=60,
    Statistics=["Average", "Maximum"],
)
for p in sorted(resp["Datapoints"], key=lambda d: d["Timestamp"]):
    print(p["Timestamp"], round(p["Average"], 1), round(p["Maximum"], 1))

Two caveats that trip people up. First, GPUUtilization on a multi-GPU instance is summed across devices, so an 8-GPU node that is perfectly busy reports ~800%, not ~100%. Divide before you panic. Second, "utilization" here means "a kernel was resident", not "the GPU was doing useful math" — a job stuck in a tiny all_gather loop can look busy at 90% and still be slow. Level 0 tells you whether you have an obvious stall, not whether you are efficient.

Read the shape of the curve:

• Sawtooth down to near zero every N seconds — the data loader is starving the GPU, or you are checkpointing synchronously.
• Flat and low the whole run — batch size too small, or you are CPU-bound in preprocessing or tokenization.
• High GPU memory, low utilization — fragmentation or a tiny per-step workload with a large cached allocator.
• First few minutes at zero — normal: container pull, S3 download, and (on newer PyTorch) compilation.

Level 1: step timing from inside the loop

The cheapest real instrumentation is a step timer that you log as a metric. Two numbers matter: seconds per step, and the fraction of each step spent waiting for data.

import time, torch

data_wait = 0.0
step_start = time.perf_counter()
t0 = time.perf_counter()

for step, batch in enumerate(loader):
    data_wait += time.perf_counter() - t0
    loss = model(**batch).loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

    if step % 50 == 0 and step:
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - step_start
        print(f"step_time_sec={elapsed / 50:.4f};")
        print(f"data_wait_frac={data_wait / elapsed:.4f};")
        print(f"tokens_per_sec={(50 * tokens_per_batch) / elapsed:.1f};")
        data_wait, step_start = 0.0, time.perf_counter()
    t0 = time.perf_counter()

The torch.cuda.synchronize() before you measure is not optional. CUDA launches are asynchronous, so without it you are timing the Python loop, not the GPU.

Wire those prints into SageMaker AI metric definitions so they become searchable CloudWatch metrics and show up in the console's training job charts:

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    entry_point="train.py",
    role=role,
    framework_version="2.6",
    py_version="py312",
    instance_type="ml.g6e.12xlarge",
    instance_count=1,
    metric_definitions=[
        {"Name": "step_time_sec", "Regex": "step_time_sec=([0-9\\.]+);"},
        {"Name": "data_wait_frac", "Regex": "data_wait_frac=([0-9\\.]+);"},
        {"Name": "tokens_per_sec", "Regex": "tokens_per_sec=([0-9\\.]+);"},
    ],
    keep_alive_period_in_seconds=1800,
)

If you are already tracking runs in managed MLflow, log the same three values per step there instead of re-deriving them from CloudWatch later. A data_wait_frac above roughly 0.1 means your input pipeline is the bottleneck, and no amount of GPU shopping will help.

Level 2: a real kernel trace with torch.profiler

When the coarse numbers say "GPU is busy but throughput is bad", you need a trace. Profile a handful of steps in the middle of training — never the first steps, which are dominated by warmup and autotuning — and write the trace to /opt/ml/output/tensorboard or straight to S3 so it survives the job.

from torch.profiler import profile, schedule, ProfilerActivity, tensorboard_trace_handler

prof = profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=20, warmup=3, active=5, repeat=1),
    on_trace_ready=tensorboard_trace_handler("/opt/ml/output/tensorboard"),
    record_shapes=True,
    profile_memory=True,
    with_stack=False,
)

prof.start()
for step, batch in enumerate(loader):
    train_step(batch)
    prof.step()
    if step > 40:
        break
prof.stop()

Keep with_stack=False for distributed runs unless you need it; Python stack capture on eight ranks can add enough overhead to change the thing you are measuring. Also profile rank 0 only (if int(os.environ.get("RANK", 0)) == 0:) or you will write eight multi-gigabyte traces into the same prefix.

Download the trace and open it in TensorBoard or at chrome://tracing / Perfetto. What to look for, in priority order:

1. Gaps on the CUDA stream with activity on the CPU thread — launch-bound. Too many tiny kernels; fix with larger batches, torch.compile, or fused optimizers.
2. Long aten::copy_ or cudaMemcpyAsync blocks — host-to-device transfer in the critical path. Pin memory, prefetch, and stop calling .item() or .cpu() inside the step.
3. nccl:all_gather / nccl:reduce_scatter dominating the timeline — communication-bound. This is the classic FSDP-on-too-many-nodes symptom; see the sharding notes below.
4. aten::to with a dtype change on every step — you are casting in the hot loop instead of loading in the target dtype.

For a summary without leaving the log, print the top kernels at the end of the profiled window:

print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=15))

The three fixes that matter most

1. Feed the GPU

On SageMaker AI, input mode is a one-line decision with large consequences. File mode copies the whole channel to local NVMe before training starts; FastFile streams objects on first read; Pipe and the S3 data source with ShardedByS3Key distribute shards across ranks.

from sagemaker.inputs import TrainingInput

train_input = TrainingInput(
    s3_data="s3://my-bucket/datasets/tokenized/",
    input_mode="FastFile",
    distribution="ShardedByS3Key",
)
estimator.fit({"train": train_input})

Rules of thumb we apply on client work: many small files plus FastFile equals per-object latency death — pack into shards of 100–500 MB (webdataset tars, Parquet, or pre-tokenized arrow files) first. If the dataset fits on the instance's NVMe and you will do multiple epochs, File mode is usually fastest overall. Then set num_workers to roughly the vCPU count divided by GPUs per node, with pin_memory=True, persistent_workers=True, and prefetch_factor=4.

Tokenizing on the fly inside the data loader is the single most common cause of a flat 30% utilization curve. Tokenize once in a SageMaker Processing job and train on the packed output.

2. Make each step do more math

• Use bf16 (torch.autocast("cuda", dtype=torch.bfloat16)) plus a bf16-friendly optimizer; skip fp16 loss scaling on Ampere and later.
• Turn on torch.compile(model) and accept the first-step penalty; on transformer blocks this often buys 10–30%.
• Enable activation checkpointing only where memory forces it — it trades roughly 30% more compute for memory, so applying it to every layer when you did not need to is a self-inflicted slowdown.
• Set torch.backends.cuda.matmul.allow_tf32 = True and use sequence lengths that are multiples of 8 (ideally 64) so tensor cores stay fed.
• Raise per-device batch size until you are at about 85% of GPU memory, then use gradient accumulation for the rest of your target global batch — accumulation steps are cheap, extra all_gather rounds are not.

3. Stop paying for communication you do not need

If the trace is full of NCCL, you have three levers. Shift from full sharding to SHARD_GRAD_OP (ZeRO-2 style) when the model fits — it removes a parameter all_gather per step. Keep the sharding group inside one node when the model fits on one node, so cross-node traffic carries gradients only. And use EFA-capable instances (ml.p5.48xlarge, ml.p4d.24xlarge, ml.trn1.32xlarge) for any multi-node run; multi-node FSDP over plain networking is a way to make eight GPUs slower than four.

Turn the finding into a guardrail

Profiling once is a consulting deliverable. Profiling continuously is an operating practice. Two cheap habits:

Alarm on idle GPUs. A CloudWatch alarm on GPUUtilization below a floor for 15 consecutive minutes catches wedged jobs, dead data loaders, and the job someone left running after the notebook crashed.

aws cloudwatch put-metric-alarm \
  --alarm-name sm-training-gpu-idle \
  --namespace /aws/sagemaker/TrainingJobs \
  --metric-name GPUUtilization \
  --statistic Average --period 300 --evaluation-periods 3 \
  --threshold 15 --comparison-operator LessThanThreshold \
  --treat-missing-data notBreaching \
  --alarm-actions arn:aws:sns:us-east-1:123456789012:ml-ops-alerts

Record cost per unit of work, not per hour. Log tokens_per_sec (or samples/sec) for every run and divide by the instance's on-demand rate. That single derived number — dollars per million tokens trained — is what tells you whether last month's "optimization" actually worked, and it is the number a finance reviewer will understand.

A worked example

A recent client engagement had an 8B-parameter fine-tune on two ml.p4d.24xlarge nodes taking 19 hours per epoch. Level 0 showed a sawtooth: 8-GPU-normalized utilization bouncing between 20% and 95% on a ~4-second period. Level 1 put data_wait_frac at 0.41. The Level 2 trace confirmed the loader was JSON-parsing and tokenizing 40,000 small S3 objects per epoch.

The fix was not a bigger instance. We pre-tokenized into 240 packed shards with a Processing job, switched the channel to FastFile with ShardedByS3Key, raised num_workers to 12, and dropped full sharding to SHARD_GRAD_OP. Epoch time went to 6.5 hours on the same two nodes — a 2.9x speedup and roughly a 65% cut in cost per epoch, with zero change to the model code.

That is the usual shape of this work: the expensive resource is the GPU-hour, and the cheapest way to buy more capacity is to stop wasting the capacity you already pay for.

---

If your training jobs are slow and you do not know why, our senior SageMaker AI consultants do this profiling work as a short, fixed-scope engagement — a trace, a ranked list of bottlenecks, and the pull request that fixes them. Get in touch with your instance types and current step times.