Training runs & jobs

Prepare and submit#

A training run combines a native Slurm batch job with experiment metadata, parameters, tags and input provenance. A run’s resources come from the declared CPU, memory, GPU and time request.

Use a unique request ID for each intentional submission. A retry with the same ID and identical arguments returns the same run. Changing arguments under an existing ID produces a conflict instead of another job.

A consistent run environment#

Keep source and inputs in /workspace. Check research_profile for the active network mode; internet-enabled allocations allow public HTTP/HTTPS and DNS. GPU allocations use the prepared ROCm/PyTorch runtime; CPU runs use the available CPU runtime. Keep credentials out of scripts, parameters and outputs.

Run on GPU compute#

Use research_profile to check your per-job limits, then research_validate to check a specific request. Request the number of GPUs your script uses, enough CPU and memory for data loading, and a realistic time limit. After submission, monitor the run’s scheduling status and start estimate when available. The GPU runtime uses PyTorch’s torch.cuda interface for AMD ROCm.

import json, os
from pathlib import Path
import torch

params = json.loads(os.environ["MATILDA_PARAMS_JSON"])
out = Path(os.environ["MATILDA_RUN_DIR"])
torch.manual_seed(7)
model = torch.nn.Linear(32, 1).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
start = 0
checkpoint = os.environ.get("MATILDA_CHECKPOINT")
if checkpoint:
    saved = torch.load(checkpoint, map_location="cuda", weights_only=True)
    model.load_state_dict(saved["model"])
    optimizer.load_state_dict(saved["optimizer"])
    start = saved["epoch"]
x = torch.randn(512, 32, device="cuda")
y = x.sum(dim=1, keepdim=True)
with (out / "metrics.jsonl").open("a") as metrics:
    for epoch in range(start, params.get("epochs", 3)):
        optimizer.zero_grad()
        loss = (model(x) - y).square().mean()
        loss.backward()
        optimizer.step()
        record = {"epoch": epoch + 1, "loss": loss.item()}
        metrics.write(json.dumps(record) + "\n")
        metrics.flush()
        print(record, flush=True)
torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(),
            "epoch": params.get("epochs", 3)}, out / "checkpoint.pt")

Save this as train.py using the workspace text tool, then submit a script that runs python3 train.py. Start with one GPU and an epochs parameter of 3. Declare the training file’s digest if you want it verified before execution.

Train on several GPUs#

Launch one process per GPU and keep --num_processes equal to the gpus in your run request. For multi-node jobs, use the platform-provided MASTER_ADDR and MASTER_PORT consistently on every rank. The example below is for one node. The workspace environment has no accelerate executable of its own; run the launcher as a module from your workspace interpreter.

export MASTER_PORT=$((RANDOM % 987 + 29000))
/workspace/.venv/bin/python -m accelerate.commands.launch \
  --num_processes=4 --num_machines=1 --dynamo_backend=no \
  --main_process_port "${MASTER_PORT}" train.py

Size the rest of the request around the GPUs: allow several CPU cores and enough memory per GPU for your data loader workers. A four-GPU vision run with four loader workers per process is comfortable at sixteen cores and 64 GiB. Call torch.distributed.destroy_process_group() on each rank at shutdown. Check distributed warnings against the exit status, logs and checkpoint; a warning alone is not evidence that all ranks completed. After the run, elapsed_seconds and allocated_gpu_seconds tell you how to size the next request.

Monitor and interpret state#

Log pages contain at most 64 KiB per read. Metrics pages consume complete newline-terminated records and report invalid lines. Keep stdout and stderr offsets separately. Poll according to poll_after_seconds instead of resubmitting work.

Resume deliberately#

After the parent run finishes, call research_resume with its run ID, the full relative workspace checkpoint path runs/<parent-id>/checkpoint.pt, a new run specification and a new UUID request ID. Use the actual parent ID in the path; a filename by itself is insufficient. The service records checkpoint SHA-256 and parent lineage, then creates a new native job.

Fine-tune a shared model and save adapters#

Get the prepared Qwen2.5 0.5B path from research_profile and put it in the run parameter model_path. Begin with one GPU, 8192 MiB of CPU memory and a short measured time limit. This example saves only the LoRA adapter and tokenizer rather than another complete base model.

import json, os
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model

params = json.loads(os.environ["MATILDA_PARAMS_JSON"])
path = params["model_path"]
out = Path(os.environ["MATILDA_RUN_DIR"])
tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=True,
                                         trust_remote_code=False)
base = AutoModelForCausalLM.from_pretrained(path, local_files_only=True,
    trust_remote_code=False, dtype=torch.bfloat16).cuda()
model = get_peft_model(base, LoraConfig(r=8, lora_alpha=16,
    target_modules=["q_proj", "v_proj"], task_type="CAUSAL_LM"))
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
batch = tokenizer("Research compute in Australia.", return_tensors="pt").to("cuda")
model.train()
with (out / "metrics.jsonl").open("w") as metrics:
    for step in range(3):
        optimizer.zero_grad()
        loss = model(**batch, labels=batch["input_ids"]).loss
        loss.backward()
        optimizer.step()
        record = {"step": step + 1, "loss": float(loss.detach().cpu())}
        metrics.write(json.dumps(record) + "\n")
        metrics.flush()
        print(record, flush=True)
model.save_pretrained(out / "adapter")
tokenizer.save_pretrained(out / "adapter")

For inference, load the same immutable base revision and attach the saved adapter using PeftModel.from_pretrained. For resumable training, also save optimizer state, progress and random state; an adapter by itself does not restore the optimizer.

Memory limits and checkpoint grace#

Allocations expose MATILDA_MEMORY_MB and MATILDA_CPUS with native Slurm limits. /proc/meminfo describes the host and is not your allowance. GPU jobs need at least 4096 MiB of CPU memory even for an import-only workload. OUT_OF_MEMORY reports the enforced limit and suggests reducing batch size or increasing the next memory request.

At a time limit or cancellation, Slurm sends SIGTERM and allows its configured KillWait grace before SIGKILL. The current grace is 15 seconds. Install a SIGTERM handler that flushes a small checkpoint promptly, and also checkpoint regularly before the deadline. An OOM kill provides no guaranteed cleanup opportunity.

import signal
from pathlib import Path
stopping = False
def stop(_signal, _frame):
    global stopping
    stopping = True
signal.signal(signal.SIGTERM, stop)
# In your training loop, check stopping between short steps. Write a temporary
# checkpoint, flush it, and atomically replace the final checkpoint path.
# Finish within the grace period; a long GPU operation can delay Python handlers.

research_run and research_job return elapsed_seconds and allocated_gpu_seconds when native accounting is available. Log reads remain bounded and expose stdout/stderr text through MCP and the HTTP API, without depending on a browser dialog.