Polyaxon & TRL

How to use Polyaxon and TRL together

+

TRL owns the training loop; Accelerate, FSDP, or DeepSpeed can own process distribution. Polyaxon packages that stack as a repeatable operation and keeps its configuration, logs, lineage, and artifacts together.

See the TRL CLI documentation for the upstream configuration and requirements.

Prerequisites

• A Kubernetes GPU node compatible with the CUDA stack in the pinned TRL image.
• A Hugging Face token connection only when the selected model or dataset is gated.
• A Polyaxon artifacts store configured for checkpoints and the trained model.

Configuration

Define the TRL component

Use TRL's published image and CLI, and expose the model and dataset as component inputs. Save the component as trl.yaml.

version: 1.1
kind: component
name: trl-sft
inputs:
- name: model
  type: str
  value: Qwen/Qwen2.5-0.5B
- name: dataset
  type: str
  value: trl-lib/Capybara
run:
  kind: job
  container:
    image: huggingface/trl:1.9.2
    command: ["trl", "sft"]
    args:
    - "--model_name_or_path"
    - "{{ model }}"
    - "--dataset_name"
    - "{{ dataset }}"
    - "--output_dir"
    - "{{ globals.run_outputs_path }}/qwen2.5-0.5b-sft"
    resources:
      limits:
        nvidia.com/gpu: "1"

This uses Hugging Face's published TRL image and the public model and dataset from its CLI example. Pin the image digest and model and dataset revisions before production.

Submit the operation

Run the component through the target Polyaxon project, queue, preset, and approval path.

polyaxon run -f trl.yaml

Compare the resolved training run

Confirm the image, revisions, parameters, accelerator allocation, logs, checkpoints, and evaluation metrics before scaling out.

polyaxon ops dashboard

Deployment checks

• Pin the container image, framework versions, model revision, dataset revision, and training configuration.
• Resolve registry, dataset, model, and tracking credentials through scoped Polyaxon connections.
• Persist checkpoints and final artifacts outside the pod, and define resume behavior before a long run.
• Set requests, limits, node selectors, tolerations, queue policy, timeout, and retry ownership explicitly.
• For multi-node training, choose one distribution owner—such as the PyTorchJob operator or a framework launcher—and validate rendezvous and checkpoint recovery.

Troubleshooting

The pod remains Pending

Inspect quota, queue capacity, node selectors, taints, topology, and the requested accelerator count.

The process runs out of memory

Check model precision, sequence length, batch size, checkpointing, optimizer state, sharding, and actual device memory.

The run cannot read or save

Verify the referenced connection, secret scope, mounted path, object-store permissions, and available storage.

References

• TRL CLI documentation — Official SFT command and example model, dataset, and output arguments.
• Hugging Face TRL container — Published image tags, including the pinned version used in this example.
• Polyaxon run outputs path — Runtime context path used to persist the trained model with the operation's outputs.
• Polyaxon job runtime — Job containers, resources, connections, initialization, and lifecycle settings.
• Polyaxon resource scheduling — Kubernetes requests, limits, custom resources, queues, and placement controls.