Run pragmatiq on Amazon SageMaker

pragmatiq is an independent implementation inspired by the PRAGMA paper (arXiv 2604.08649) and is not affiliated with or endorsed by Revolut.

This guide runs pragmatiq entirely on Amazon SageMaker — no servers to manage. You prep data in SageMaker Studio, train with a managed SageMaker training job, and serve embeddings from a NVIDIA Triton real-time endpoint. pragmatiq is on PyPI, so every step is just pip install pragmatiq.

It mirrors the local Pretrain and Serve with Triton tutorials — same commands, same Triton request shape — on SageMaker-managed infrastructure. Works on the synthetic generator or on your own bank data (same four-file contract).

How pragmatiq works — a foundation model for events, not text

A large language model reads a stream of subword tokens and predicts the next one. pragmatiq reads a stream of timestamped key–value events — a card payment, an app screen, a transfer — and learns to reconstruct masked fields. It's the same transformer backbone, but because the input isn't text, three things change: the tokenizer, the notion of position, and the training objective.

Click through the stack — every stage is implemented in pragmatiq/models/:

The tokenizer: key–value–time, not subwords

A text tokenizer has one job: split a string into subword ids. pragmatiq's tokenizer fits one vocabulary over the whole dataset and turns each field into a key token plus a value representation chosen by the field's kind — keys and values share one embedding space:

• Numeric values (amounts, balances) are percentile-binned with a dedicated zero bucket, so "no balance" is its own symbol and an unseen magnitude clips into the end buckets instead of failing.
• Low-cardinality strings (country, channel) become one categorical token per value.
• High-cardinality text (merchant names, device ids) is byte-level BPE by default — or, in the Nemotron variant, a single sentinel token whose raw string a frozen text model maps to a vector.

Unseen keys or values at inference map to [UNK] with a warning — the model never raises on vocabulary drift, which matters when you serve a live book.

Time is the spine, not an afterthought

Banking events are wildly irregular — seconds apart, then months apart — so an integer position ("the previous event") throws away the strongest signal. pragmatiq compresses elapsed seconds with a log curve:

and feeds that as a continuous position for rotary embeddings (TimeRoPE) — a token at log-seconds $p$ rotates frequency pair $i$ by $p \cdot \text{inv\_freq}_i$, so attention encodes the elapsed time between two events rather than their ordinal distance. Calendar features (hour, day-of-week, day-of-month) take a separate sin/cos → MLP path, so day/night and weekday/weekend structure is available independently of elapsed time.

Why the encoder stack differs from a vanilla transformer

Rather than one flat sequence, pragmatiq encodes a hierarchy so an event's fields don't bleed across event boundaries:

1. The event encoder encodes each event independently (block-diagonal attention within an event, a prepended [EVT] marker); its output plus calendar features is the per-event vector.
2. The profile-state encoder encodes static attributes + lifelong milestones under a [USR] marker.
3. The history encoder runs over [profile, event…] with TimeRoPE on log-seconds-to-the-latest-event; the [USR] slot output is the user embedding.

All three are bidirectional, pre-norm, GELU, ffn = 4d — and the whole batch is packed without padding via cu_seqlens (flash-attn on GPU, an SDPA fallback on CPU, checked against a padded reference). Pretraining masks 15% of tokens, 10% of whole events, and 10% of (user, key) groups, then reconstructs each masked value from a 3-view head — concat[ẑ_e, z_h(event), z_h(USR)] ∈ ℝ^{3d} → Linear(3d→d) — scored against the tied embedding table.

The Nemotron variant

By default, high-cardinality text is BPE. The optional PRAGMA+Nemotron variant instead emits one sentinel token per text field and lets a frozen text model embed its raw string; masked text tokens are reconstructed with MSE against that vector, so the loss becomes $\mathcal{L} = \text{CE} + \lambda \cdot \text{MSE}$. It's switchable from the data step alone and off by default (the BPE path stays byte-identical) — tokenize with configs/data/tokenizer_nemotron.yaml and pip install "pragmatiq[extras]" for the frozen embedder. See the variant page.

Set up SageMaker

aws s3 mb s3://YOUR-pragmatiq-bucket

pip install pragmatiq sagemaker

import sagemaker
sess = sagemaker.Session()
role = sagemaker.get_execution_role()
bucket = "YOUR-pragmatiq-bucket"

Get your data into S3

pragmatiq trains on a four-file parquet contract — the same files whether they're synthetic or yours. Produce them, then upload to S3.

from pragmatiq import api
api.synthesize({"n_users": 50000, "seed": 0}, out="data", n_workers=8)
sess.upload_data("data", bucket=bucket, key_prefix="data")   # -> s3://.../data/

import pyarrow as pa, pyarrow.parquet as pq

events = pa.table({
    "user_id": ["u1", "u1"],
    "ts": pa.array([1_700_000_000_000_000, 1_700_003_600_000_000], type=pa.timestamp("us")),
    "source": ["transaction", "app"],            # groups schemas: transaction / app / trading …
    "fields": [                                   # map<string,string> per event
        {"amount": "42.10", "mcc": "5411", "merchant": "TESCO STORES 4521"},
        {"screen": "home", "action": "view"},
    ],
})
pq.write_table(events, "data/events.parquet")
# + profiles.parquet (user_id, as_of, attributes, lifelong),
#   optional transfers.parquet and labels/<task>.parquet — see the data contract.

import subprocess
subprocess.run(["pragmatiq", "validate", "data"], check=True)   # fails loudly on dtype/integrity issues
sess.upload_data("data", bucket=bucket, key_prefix="data")

Train with a SageMaker training job

A SageMaker training job runs your script on a GPU instance it provisions and tears down for you, reading the data from S3 and writing the trained run back to S3. Two small files define it.

# src/train.py
import os
from pragmatiq import api

data = os.environ["SM_CHANNEL_DATA"]      # S3 "data" channel, mounted by SageMaker
out  = os.environ["SM_MODEL_DIR"]         # /opt/ml/model -> tarred to S3 as model.tar.gz

api.tokenize(data, "/tmp/tok", n_workers=8)
api.pretrain("/tmp/tok", "run", model_size="small",
             config={"max_steps": 4000, "token_budget": 16384}, runs_root=out)

pragmatiq[serve]

from sagemaker.pytorch import PyTorch

est = PyTorch(
    entry_point="train.py", source_dir="src",
    role=role, framework_version="2.4", py_version="py311",
    instance_type="ml.g5.2xlarge", instance_count=1,
)
est.fit({"data": f"s3://{bucket}/data"})
print(est.model_data)        # s3://.../model.tar.gz — the trained run

Serve with NVIDIA Triton on a SageMaker endpoint

pragmatiq's production serving path is a Triton python backend running the native varlen PyTorch model — the exact no-padding forward used in training (see Serve with Triton). SageMaker hosts it on a managed real-time endpoint: you supply a Triton model repository and a container with pragmatiq installed.

# Dockerfile — base on the current SageMaker Triton image for your region (see the
# SageMaker Triton docs for the URI), then add pragmatiq from PyPI.
FROM <ACCOUNT>.dkr.ecr.<REGION>.amazonaws.com/sagemaker-tritonserver:<TAG>
RUN pip install --no-cache-dir "pragmatiq[serve]"

aws ecr create-repository --repository-name pragmatiq-triton
# docker build -t pragmatiq-triton . && tag + push to ECR (see the ECR docs)

base=https://raw.githubusercontent.com/dynamiq-ai/pragmatiq/main/deploy/triton/model_repository/pragmatiq_embedder
mkdir -p model_repository/pragmatiq_embedder/1/run
curl -sL $base/config.pbtxt -o model_repository/pragmatiq_embedder/config.pbtxt
curl -sL $base/1/model.py   -o model_repository/pragmatiq_embedder/1/model.py

aws s3 cp s3://YOUR-pragmatiq-bucket/.../model.tar.gz run.tgz   # est.model_data
tar -xzf run.tgz -C model_repository/pragmatiq_embedder/1/run --strip-components=1

tar -C model_repository -czf triton.tar.gz pragmatiq_embedder
aws s3 cp triton.tar.gz s3://YOUR-pragmatiq-bucket/triton/triton.tar.gz

aws sagemaker create-model --model-name pragmatiq-embedder \
  --execution-role-arn arn:aws:iam::<ACCOUNT>:role/<SageMakerExecutionRole> \
  --primary-container 'Image=<ACCOUNT>.dkr.ecr.<REGION>.amazonaws.com/pragmatiq-triton:latest,ModelDataUrl=s3://YOUR-pragmatiq-bucket/triton/triton.tar.gz,Environment={SAGEMAKER_TRITON_DEFAULT_MODEL_NAME=pragmatiq_embedder,PRAGMATIQ_RUN=/opt/ml/model/pragmatiq_embedder/1/run,SAGEMAKER_TRITON_SHM_DEFAULT_BYTE_SIZE=1073741824}'

aws sagemaker create-endpoint-config --endpoint-config-name pragmatiq-embedder \
  --production-variants VariantName=main,ModelName=pragmatiq-embedder,InstanceType=ml.g5.xlarge,InitialInstanceCount=1

aws sagemaker create-endpoint --endpoint-name pragmatiq-embedder \
  --endpoint-config-name pragmatiq-embedder
aws sagemaker wait endpoint-in-service --endpoint-name pragmatiq-embedder

aws sagemaker-runtime invoke-endpoint --endpoint-name pragmatiq-embedder \
  --content-type application/json \
  --body '{"inputs":[{"name":"records_json","shape":[1],"datatype":"BYTES","data":["[{\"user_id\":\"u1\",\"events\":[{\"ts\":1718200000000000,\"source\":\"transaction\",\"fields\":{\"amount\":\"42.50\",\"merchant\":\"TESCO\"}}],\"attributes\":{\"country\":\"GB\"},\"lifelong\":[]}]"]}]}' \
  out.json && cat out.json

Clean up

A real-time endpoint bills per instance-hour while it's in service — delete it when you're done:

aws sagemaker delete-endpoint --endpoint-name pragmatiq-embedder
aws sagemaker delete-endpoint-config --endpoint-config-name pragmatiq-embedder
aws sagemaker delete-model --model-name pragmatiq-embedder

Training jobs bill only while running and stop on their own. Artifacts in S3 are cheap; delete them when finished.

Where to go next

• Amazon SageMaker training jobs and SageMaker Studio.
• NVIDIA Triton on SageMaker — the container, model repository, and autoscaling.
• The same pipeline locally: Quickstart → Pretrain → Serve, and Bring your own data.