@tuned-tensor/cli 0.15.0 → 0.16.0

package/docs/local-runtime/architecture.md

@@ -127,17 +127,28 @@ hide a valid artifact. Serving verifies the manifest again.
 
 ## Configuration
 
-The runner config controls only:
+The runner config controls:
 
 - state, artifact, and Hugging Face cache paths;
+- optional AWS instance lookup, SSH access, and a remote process deadline;
 - CUDA or CPU evaluation, with local serving through the CUDA vLLM runtime;
 - deterministic generation and scoring limits.
 
-The training project, Python entrypoints, working directory, and child
-environment are internal and fixed. Every stage uses the locked `uv` project
-included in the npm package. The mutable virtual environment lives in
-`~/.tuned-tensor/cache`, keyed by the project and lockfile content, so a
-read-only global npm install remains runnable.
+The training project, Python entrypoints, working directory, and forwarded
+runtime settings are internal and fixed. Every stage uses the locked `uv`
+project included in the npm package. Locally, the mutable virtual environment
+lives in `~/.tuned-tensor/cache`, keyed by the project and lockfile content, so
+a read-only global npm install remains runnable.
+
+With [`gpu` configured](aws-gpu.md), the same orchestrator dispatches adapter
+training and CUDA inference, or foundation stages after tokenization, over SSH
+to an existing EC2 instance. It stages the bundled Python source, lockfile and
+declared inputs in a private remote directory, creates the runtime environment
+there, and retrieves outputs before cleanup. Dataset compilation, adapter
+scoring/comparison, artifact validation and run state remain local. Foundation
+tokenization also stays local. Serving and `tt hardware` inspect/use the local
+host. AWS credentials authorize instance lookup locally; they are not forwarded
+to GPU processes. TT does not provision or manage instance capacity.
 
 Schemas are strict. Misspelled or obsolete fields fail validation instead of
 silently reverting to defaults. Training itself is CUDA-only and fails fast

package/docs/local-runtime/aws-gpu.md

@@ -0,0 +1,131 @@
+# Use your AWS GPU with the local pipeline
+
+TT runs one pipeline on your laptop. Add `gpu` to `local-runner.json` to send
+adapter training and CUDA evaluation, or foundation pretraining, fine-tuning,
+RL and evaluation, to an existing EC2 instance in your account. Foundation
+tokenization, dataset preparation, adapter scoring/comparisons, run history,
+and artifact validation stay local. Results return to the usual local
+directories. The laptop does not need an NVIDIA GPU.
+
+Keep the existing spec and pipeline. Save this config beside `tunedtensor.json`,
+replacing the example values with your own. You can also copy
+[`examples/local-runtime/aws-runner.json`](../../examples/local-runtime/aws-runner.json).
+All runner paths refer to the laptop; TT translates the paths needed remotely.
+
+```json
+{
+  "gpu": {
+    "provider": "aws",
+    "profile": "research",
+    "region": "eu-west-1",
+    "instanceId": "i-0123456789abcdef0",
+    "user": "ubuntu",
+    "identityFile": "~/.ssh/research-gpu.pem",
+    "maxSeconds": 86400
+  }
+}
+```
+
+The laptop needs AWS CLI v2, OpenSSH, rsync 3.2+, and uv. On macOS, install
+a current rsync (for example with Homebrew); the system version is too old.
+Each machine must find its required programs through `PATH`; on the instance,
+this includes noninteractive SSH sessions. The instance needs Linux, a compatible NVIDIA
+GPU/driver, uv, rsync 3.2+, bash, setsid, and GNU timeout. TT
+copies the bundled Python source and lockfile and runs that exact runtime
+with `uv run --frozen`; there is no separate training image or Step Functions
+pipeline. Internet access on the instance is needed to install the locked
+runtime. No TT token is needed for training.
+
+`profile` and `region` are optional. AWS CLI uses its normal credential chain,
+including environment credentials, configured profiles and SSO. For SSO, log
+in first with `aws sso login --profile research`. The selected identity needs
+`ec2:DescribeInstances`. TT resolves the instance using
+[AWS DescribeInstances](https://docs.aws.amazon.com/cli/latest/reference/ec2/describe-instances.html).
+An AWS profile authorizes the lookup; SSH access to the machine is also required.
+The selected instance must already be running in the selected region and have
+enough GPU memory and free disk for the model, runtime, inputs and outputs.
+
+Verify the instance's SSH host key for the selected IP address and connect
+once yourself before using TT. TT requires a known host key and noninteractive
+SSH authentication. Omit
+`identityFile` to use your SSH agent or SSH configuration. Relative identity
+file paths resolve beside the runner config; `~` expands locally. Set
+`privateIp: true` when connected to the instance's VPC through a VPN or other
+private network. The default uses its public IPv4 address. Security groups and
+network routes must allow the laptop to reach that address over SSH.
+TT addresses the instance by IP, so any `Host` rules in your SSH configuration
+must match that IP. Remote Python commands run through `bash -lc`, which loads
+the login profile so a normal uv installation can be found.
+
+```bash
+# Adapter workflow: fetch the base model locally first.
+tt validate tunedtensor.json
+tt models prefetch tunedtensor.json --config local-runner.json
+tt doctor tunedtensor.json --config local-runner.json
+tt pipeline run --spec tunedtensor.json --config local-runner.json --dry-run
+tt pipeline run --spec tunedtensor.json --config local-runner.json
+
+# Foundation workflow: no base-model prefetch is needed.
+tt doctor tunedtensor.json --config local-runner.json
+tt pipeline run --spec tunedtensor.json --config local-runner.json \
+  --resume /absolute/path/to/foundation-run
+```
+
+Choose the command block for your spec's engine. `tt doctor` connects to the
+instance and prepares/probes its locked Python/CUDA runtime, which may take
+time on first use. It does not start training or reserve GPU capacity.
+`tt hardware` reports the laptop's hardware; it does not inspect the AWS
+instance. `tt pipeline run` discovers `local-runner.json` beside the selected
+spec when `--config` is omitted.
+
+Keep pipeline step targets set to `local`: the target describes the
+orchestrator, and `gpu` selects where its GPU processes execute. CPU adapter
+evaluation (`evaluation.inference.device: "cpu"`) stays on the laptop.
+Remove `gpu` to execute everything locally. `--dry-run` or `dryRun: true`
+does not connect to AWS or copy files.
+
+TT copies only each process's declared data, model snapshot, and runtime files,
+not the workspace or laptop credentials. Adapter runs use an immutable local
+model snapshot; the Hugging Face credential/cache root is never uploaded.
+GPU processes inherit the remote machine's environment plus a small set of
+explicit runtime settings. Gated-model authentication happens during local
+prefetch. Logs stream to the laptop; models and foundation checkpoints are
+downloaded when the process ends, before temporary remote files are removed.
+Staging paths are stable for checkpoint resume and reserved exclusively; an
+existing directory blocks a new process until its previous work or artifacts
+are recovered. Inputs and runtime files are staged for each process, so large
+models incur transfer time and require space on both
+machines, including the instance's `/tmp`. Serving remains local in this version.
+
+Foundation `checkpoint_backup_dir` is also a local path that TT stages remotely.
+Keep it outside the foundation run directory; overlapping input, output,
+recovery or backup paths are rejected before transfer.
+During pretraining, both rolling checkpoints and this backup are written under
+the instance's staging directory. They return to the laptop only after the
+process ends or is cancelled; there is no continuous backup to the laptop or
+to a separate remote disk. Protect and recover that staging directory if the
+connection fails. See [long foundation runs](foundation-long-runs.md) for the
+checkpoint contents and resume contract.
+
+Cancellation signals the remote process group and retrieves available
+outputs while SSH remains reachable. If SSH or artifact retrieval fails, the
+error includes the remote staging directory for recovery; keep it until you have recovered
+your files. A remote timeout (24 hours by default, configurable from 60 seconds
+to 7 days per process) limits work after a lost connection or laptop failure.
+Adapter evaluation also retains `evaluation.timeoutMs` (30 minutes by default),
+and individual transfer/setup commands have a 30-minute timeout. Increasing
+`gpu.maxSeconds` does not extend those limits.
+The pipeline needs the laptop to stay running; it is not a detached scheduler.
+Foundation pretraining can resume from returned checkpoints using its existing
+`--resume` command. Use the same laptop run directory and unchanged spec/data
+so checkpoint identities still match. Adapter training keeps its existing
+restart behavior.
+
+You provide capacity and pay AWS directly. TT does not create, start, stop,
+terminate or resize instances, request quotas, choose fallback instances or
+fall back to another account. An AWS profile alone cannot create capacity.
+The instance remains running after the job; its lifecycle is yours to manage.
+
+Hosted training submission (`tt cloud runs start` and `estimate`) is retired.
+Existing cloud reports and cancellation of older hosted jobs remain available.
+This CLI change does not delete deployed AWS infrastructure or historical data.

package/docs/local-runtime/foundation-long-runs.md

@@ -4,6 +4,10 @@ The foundation engine can run recoverable, single-GPU pretraining from `.txt`
 files or JSONL shards whose rows contain a string `text` field. Keep training
 and validation data in different files or directories.
 
+The commands below also support an [AWS GPU](aws-gpu.md) configured in the
+adjacent `local-runner.json`. The CLI and tokenizer still run locally; the
+local machine must stay awake and connected throughout the pipeline.
+
 ## Spec
 
 The long-run fields live under `foundation` in `tunedtensor.json`:
@@ -69,11 +73,22 @@ SIGINT or SIGTERM requests a step-boundary checkpoint. The CLI allows up to two
 minutes for that checkpoint before forcibly stopping the process. SIGKILL or a
 power loss falls back to the most recent periodic checkpoint.
 
+With AWS execution, checkpoints remain in remote staging during pretraining
+and are downloaded when the process ends or cancellation succeeds. The
+configured `checkpoint_backup_dir` is mapped into that same remote staging
+tree, then copied back to its local destination; it is not a live backup to
+another disk or to the laptop. Keep the backup directory outside the run
+directory; AWS execution rejects overlapping paths. An interrupted connection
+may need manual recovery from the remote path printed in the error before
+`--resume` can run.
+Set `gpu.maxSeconds` to cover the intended pretrain stage (default 24 hours,
+maximum 7 days); reaching the deadline ends that process.
+
 ## User service
 
-For unattended operation, run the resumable command under a user systemd
-service. Replace every path with an absolute path and use the installed `tt`
-path reported by `command -v tt`:
+For unattended operation on a Linux machine running the CLI, run the resumable
+command under a user systemd service. Replace every path with an absolute path
+and use the installed `tt` path reported by `command -v tt`:
 
 ```ini
 [Unit]
@@ -102,7 +117,8 @@ journalctl --user -fu tt-foundation.service
 ```
 
 Enable user lingering if the job must survive logout. Use a UPS for power-loss
-protection, and put `checkpoint_backup_dir` on a different physical disk or
-remote mounted filesystem. Before a multi-day run, complete a full-load soak,
+protection. With a local GPU, put `checkpoint_backup_dir` on a different
+physical disk or remote mounted filesystem. AWS staging does not provide that
+separate-disk guarantee. Before a multi-day run, complete a full-load soak,
 stop the service during pretraining, and verify the journal reports a resumed
 checkpoint after restart.

package/docs/local-runtime/serving.md

@@ -13,11 +13,21 @@ not mutate the training runtime.
   serving is not supported by this runtime; CLI/package commands still work
   there. Unsupported hardware or model/adapter combinations fail rather than
   silently falling back to the old server.
+- `gpu` in `local-runner.json` applies to training and evaluation. It does not
+  move `tt serve` to AWS; returned artifacts still need a Linux/CUDA machine
+  running the serving command.
 - Install `uv` and prefetch the certified model with `tt models prefetch`.
   First launch installs the locked serving dependencies and may compile kernels;
   it needs substantially more time and disk than a warm request. CUDA JIT may
   require a working CUDA toolkit/C++ toolchain. Model loading remains offline.
-- Foundation checkpoints and image inputs are not supported by this path.
+- Foundation model directories are supported through a temporary GPT-2 export;
+  pass `tt serve foundation --checkpoint <model-dir> --tokenizer <tokenizer.json>`.
+  The tokenizer must be the one used during training. Export checks vocabulary
+  size and required special tokens, not tokenizer identity: a different tokenizer
+  with the same vocabulary size can pass these checks and produce incorrect text.
+  Keep each model paired with its original training tokenizer. Context defaults to
+  the checkpoint limit. Text completions and chat are supported; foundation tool-call
+  parsing and image inputs are not supported.
 - Qwen3.5-2B base and a one-step, rank-16 `all-linear` TT-shaped adapter have been
   exercised on GB10. The registry's Nemotron and Muse Glimmer architectures have
   upstream support, but their exact adapters/hardware combinations are **not

package/examples/local-runtime/aws-runner.json

@@ -0,0 +1,13 @@
+{
+  "artifactRoot": ".tuned-tensor/artifacts",
+  "storeRoot": ".tuned-tensor/store",
+  "gpu": {
+    "provider": "aws",
+    "profile": "research",
+    "region": "eu-west-1",
+    "instanceId": "i-0123456789abcdef0",
+    "user": "ubuntu",
+    "identityFile": "~/.ssh/research-gpu.pem",
+    "maxSeconds": 86400
+  }
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tuned-tensor/cli",
-  "version": "0.15.0",
+  "version": "0.16.0",
   "description": "Local TT agent and CUDA fine-tuning CLI for Tuned Tensor workflows",
   "repository": {
     "type": "git",
@@ -13,12 +13,15 @@
   "scripts": {
     "build": "tsup",
     "dev": "tsup --watch",
-    "typecheck": "tsc --noEmit",
+    "typecheck": "tsc --noEmit && tsc --noEmit -p tsconfig.scripts.json",
     "test": "npm run test:cli && npm run test:local-runtime && npm run test:python",
     "test:cli": "vitest run",
+    "test:workflows": "vitest run src/__tests__/workflows",
+    "eval:agent": "tsx scripts/eval-agent.ts",
     "test:local-runtime": "node --import tsx --test \"test/local-runtime/**/*.test.ts\"",
-    "test:python": "python3 -m unittest discover -s training/adapter/tests -p 'test_*.py' && uv run --frozen --project training/foundation python -m unittest discover -s training/foundation/tests -p 'test_*.py'",
-    "test:watch": "vitest"
+    "test:python": "python3 -m unittest discover -s training/adapter/tests -p 'test_*.py' && uv run --frozen --project training/foundation --group test python -m unittest discover -s training/foundation/tests -p 'test_*.py'",
+    "test:watch": "vitest",
+    "check": "npm run typecheck && npm test && npm run build"
   },
   "files": [
     "dist",
@@ -50,9 +53,9 @@
     "@earendil-works/pi-agent-core": "0.84.1",
     "@earendil-works/pi-coding-agent": "0.84.1",
     "@tuned-tensor/pipeline-contract": "1.1.0",
-    "chalk": "^5.6.2",
+    "chalk": "^6.0.0",
     "cli-table3": "^0.6.5",
-    "commander": "^14.0.3",
+    "commander": "^15.0.0",
     "open": "^10.1.0",
     "ora": "^9.3.0",
     "typebox": "1.3.7",

package/training/adapter/src/foundation_export.py

@@ -0,0 +1,81 @@
+"""Export TT FoundationGPT weights to vLLM's existing GPT-2 implementation."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+CHAT_TEMPLATE = """{%- if not messages or messages[0]['role'] != 'system' -%}
+{{- '<|system|>You are a helpful assistant.<|end|>' -}}
+{%- endif -%}
+{%- for message in messages -%}
+{%- set content = message['content'] -%}
+{%- if message['role'] == 'system' -%}
+{%- set content = (content | trim) or 'You are a helpful assistant.' -%}
+{%- endif -%}
+{%- if message['role'] not in ['system', 'user', 'assistant'] -%}
+{{- raise_exception('Foundation chat supports system, user and assistant messages only.') -}}
+{%- endif -%}
+{{- '<|' + message['role'] + '|>' + content + '<|end|>' -}}
+{%- endfor -%}
+{%- if add_generation_prompt -%}{{- '<|assistant|>' -}}{%- endif -%}"""
+
+
+def export_foundation(checkpoint: str, tokenizer_path: str, destination: Path) -> str:
+    import torch
+    from safetensors.torch import load_file, save_file
+    from tokenizers import Tokenizer
+    from transformers import GPT2Config, PreTrainedTokenizerFast
+
+    source = Path(checkpoint)
+    config = json.loads((source / 'config.json').read_text())
+    for key in ('depth', 'width', 'heads', 'vocab_size', 'sequence_length'):
+        if type(config.get(key)) is not int or config[key] <= 0:
+            raise ValueError(f'Invalid foundation config: {key}')
+    if config['width'] % config['heads']:
+        raise ValueError('Foundation width must be divisible by heads.')
+    raw_tokenizer = Tokenizer.from_file(tokenizer_path)
+    special_tokens = ['<unk>', '<pad>', '<|system|>', '<|user|>', '<|assistant|>', '<|end|>']
+    if any(raw_tokenizer.token_to_id(token) is None for token in special_tokens):
+        raise ValueError('Foundation tokenizer is missing required special tokens.')
+    tokenizer = PreTrainedTokenizerFast(tokenizer_object=raw_tokenizer,
+        unk_token='<unk>', pad_token='<pad>', eos_token='<|end|>',
+        additional_special_tokens=['<|system|>', '<|user|>', '<|assistant|>'],
+        model_max_length=config['sequence_length'], chat_template=CHAT_TEMPLATE)
+    if len(tokenizer) != config['vocab_size']:
+        raise ValueError('Foundation tokenizer vocabulary does not match the checkpoint.')
+    destination.mkdir(parents=True, exist_ok=True)
+    tokenizer.save_pretrained(destination)
+    GPT2Config(n_layer=config['depth'], n_embd=config['width'], n_head=config['heads'],
+        n_positions=config['sequence_length'], vocab_size=config['vocab_size'],
+        activation_function='gelu', layer_norm_epsilon=1e-5,
+        resid_pdrop=0.0, embd_pdrop=0.0, attn_pdrop=0.0,
+        bos_token_id=None, eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id, architectures=['GPT2LMHeadModel'],
+        tie_word_embeddings=True).save_pretrained(destination)
+    weights = load_file(str(source / 'model.safetensors'))
+    exported = {}
+    def move(old, new, transpose=False):
+        tensor = weights.pop(old)
+        exported[new] = tensor.t().contiguous() if transpose else tensor
+    move('tok_emb.weight', 'transformer.wte.weight')
+    move('pos_emb.weight', 'transformer.wpe.weight')
+    head = weights.pop('lm_head.weight')
+    if not torch.equal(head, exported['transformer.wte.weight']):
+        raise ValueError('Foundation checkpoint must have tied embedding weights.')
+    for suffix in ('weight', 'bias'):
+        move('ln_f.' + suffix, 'transformer.ln_f.' + suffix)
+    for i in range(config['depth']):
+        old, new = f'blocks.{i}.', f'transformer.h.{i}.'
+        for src, dst in [('ln1', 'ln_1'), ('ln2', 'ln_2'),
+                         ('mlp.0', 'mlp.c_fc'), ('mlp.2', 'mlp.c_proj')]:
+            for suffix in ('weight', 'bias'):
+                move(old + src + '.' + suffix, new + dst + '.' + suffix,
+                     transpose=src.startswith('mlp') and suffix == 'weight')
+        for src, dst in [('attn.qkv', 'attn.c_attn'), ('attn.proj', 'attn.c_proj')]:
+            move(old + src + '.weight', new + dst + '.weight', transpose=True)
+            weight = exported[new + dst + '.weight']
+            exported[new + dst + '.bias'] = torch.zeros(weight.shape[1], dtype=weight.dtype)
+    if weights:
+        raise ValueError(f'Unexpected foundation weights: {sorted(weights)}')
+    save_file(exported, str(destination / 'model.safetensors'), metadata={'format': 'pt'})
+    return str(destination)

package/training/adapter/src/serve.py

@@ -64,9 +64,10 @@ def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path)
         "Qwen/Qwen3.5-2B": "qwen3_xml",
         "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16": "qwen3_coder",
         "meta-models/Muse-Glimmer-30B": "muse_glimmer",
-    }[base]
-    args += ["--enable-auto-tool-choice", "--tool-call-parser", parser]
-    if base != "meta-models/Muse-Glimmer-30B":
+    }.get(base)
+    if parser:
+        args += ["--enable-auto-tool-choice", "--tool-call-parser", parser]
+    if parser and base != "meta-models/Muse-Glimmer-30B":
         args += ["--default-chat-template-kwargs", '{"enable_thinking":false}']
     if adapter_path:
         config = json.loads((Path(adapter_path) / "adapter_config.json").read_text())
@@ -124,11 +125,17 @@ def run_server() -> None:
     family = socket.AF_INET6 if ":" in host else socket.AF_INET
     with socket.socket(family, socket.SOCK_STREAM) as probe:
         probe.bind((host, port))
-    model_source = prepare_model_source()
     if os.environ.get("TT_API_KEY"):
         os.environ["VLLM_API_KEY"] = os.environ["TT_API_KEY"]
     with TemporaryDirectory(prefix="tt-serve-") as temp:
-        adapter = resolve_adapter_path(os.environ.get("TT_MODEL_ARTIFACT"), Path(temp))
+        if os.environ.get("TT_FOUNDATION_CHECKPOINT"):
+            from foundation_export import export_foundation
+            model_source = export_foundation(os.environ["TT_FOUNDATION_CHECKPOINT"],
+                os.environ["TT_FOUNDATION_TOKENIZER"], Path(temp) / "foundation")
+            adapter = None
+        else:
+            model_source = prepare_model_source()
+            adapter = resolve_adapter_path(os.environ.get("TT_MODEL_ARTIFACT"), Path(temp))
         args = build_vllm_args(model_source, adapter, Path(temp))
         from vllm.entrypoints.cli.main import main
         previous = sys.argv

package/training/adapter/src/train.py

@@ -282,6 +282,10 @@ def run_training(rows: list[dict[str, Any]], model_source: str) -> tuple[dict[st
         )
         result = trainer.train()
 
+    # Saved adapters must remain loadable after a local or remote staging directory is removed.
+    for adapter_config in model.peft_config.values():
+        adapter_config.base_model_name_or_path = str(hp("base_model", CERTIFIED_BASE_MODEL))
+        adapter_config.revision = hp("base_model_revision")
     model.save_pretrained(MODEL_DIR)
     tokenizer.save_pretrained(MODEL_DIR)
     return result.metrics, dtype

package/training/foundation/pyproject.toml

@@ -10,6 +10,10 @@ dependencies = [
   "torch>=2.11.0",
 ]
 
+[dependency-groups]
+# Match the pinned serving runtime for CPU export parity tests.
+test = ["transformers==5.16.1"]
+
 [tool.uv.sources]
 torch = [
   { index = "pytorch-cu130", marker = "sys_platform == 'linux'" },

package/training/foundation/src/model.py

@@ -132,6 +132,8 @@ def generate(
 def save_model(model: FoundationGPT, directory: str | Path) -> None:
     destination = ensure_private_directory(directory)
     tensors = {key: value.detach().cpu().contiguous() for key, value in model.state_dict().items()}
+    # CPU tensors can still alias the tied embedding after .cpu().
+    tensors["lm_head.weight"] = tensors["lm_head.weight"].clone()
     model_path = destination / "model.safetensors"
     save_file(tensors, str(model_path))
     make_private_file(model_path)