Schreiben Sie Ihre eigene Trainingsschleife — ohne GPUs verwalten zu müssen — mit den Amazon SageMaker Training Session APIs

Das Feintuning eines großen Sprachmodells (LLM) erfordert heute meist einen Kompromiss zwischen Kontrolle über die Trainingsschleife und operativer Einfachheit. Bei vollständig verwalteten Services wählen Sie ein Basismodell aus, liefern Daten und Hyperparameter, und der Service übernimmt alles andere – das funktioniert gut, wenn ein standardisiertes, vorgefertigtes Rezept zu Ihrem Problem passt. Weniger verwaltete Services geben Ihnen mehr Kontrolle, aber Sie verwalten dann auch mehr der zugrunde liegenden Infrastruktur und Trainings-Frameworks. Manchmal reicht ein Standardrezept nicht aus: Sie benötigen eine eigene Loss-Funktion oder Daten-Curriculum für Supervised Fine-Tuning (SFT), eine eigene Präferenzoptimierung, eine Reinforcement-Learning-(RL-)Reward-Funktion oder einen Alignment-Algorithmus, der letzte Woche veröffentlicht wurde. Sobald Sie die Trainingsprozedur selbst gestalten müssen, übernehmen Sie auch diese Infrastrukturarbeit, um algorithmische Freiheit zu gewinnen.

Heute wird das serverlose Modell-Customizing in Amazon SageMaker AI um Training Session APIs erweitert, die diesen Kompromiss beseitigen, indem sie die Abstraktionsgrenze eine Ebene tiefer ziehen. Statt den Algorithmus zu abstrahieren, abstrahieren sie die Compute-Ressourcen. Operationen wie forward_backward(), optim_step(), und sample() werden zu SageMaker-API-Aufrufen, die Sie zu der Prozedur zusammen setzen, die Ihre Technik erfordert: SFT, Präferenzoptimierung, RL oder etwas Neues. Die Aufrufe zielen auf eine aktive, zustandsbehaftete Sitzung, sodass das Modell über mehrere Schritte hinweg auf den GPUs geladen bleibt, und Sie können so viele Aufrufe dagegen ausführen, wie Sie benötigen. Der Service übernimmt Runtime-Setup, verteilte Ausführung, Gradientensynchronisierung, Fehlerbehebung, Tenant-Isolation und die Operationen des GPU-Betriebs im großen Maßstab. Sie schreiben die Trainingsschleife, und SageMaker AI betreibt die Infrastruktur.

Zusammen mit dem Service veröffentlichen wir das SageMaker training sessions cookbook, eine Sammlung sofort ausführbarer SFT- und RL-Rezepte, die Sie als Blaupausen für Ihre eigenen Trainingsschleifen verwenden können.

In diesem Beitrag stellen wir das Programmiermodell und die drei SDK-Clients vor und zeigen, wo die Training Session APIs unter den Modelltrainings- und Customizing-Optionen in SageMaker AI einzuordnen sind. Anschließend zerlegen wir eine minimale RL-Schleife für Mathematikaufgaben der Grundschulstufe (GSM8K), um zu zeigen, wie die Operationen zusammen wirken, und schließen mit einer realen Workload aus dem Kochbuch: Reinforcement Learning auf Competitive-Programming-Aufgaben, bei denen generierter Code bewertet wird, indem er gegen Unit-Tests ausgeführt wird.

Was sind die SageMaker AI Training Session APIs?

Mit den SageMaker AI Training Session APIs können Sie eine aktive, GPU-gestützte Sitzung aus Ihrem eigenen Python-Prozess auf Ihrem lokalen Rechner steuern. Sie rufen die Operationen auf, und SageMaker AI führt sie auf verwalteten Beschleunigern aus, auf denen Ihr gewähltes Basismodell plus Ihr Low-Rank Adaptation (LoRA) Adapter geladen ist. Es gibt keine Konfigurationsdatei, die den gesamten Job im Voraus beschreibt, und keinen Schritt zur Cluster-Erstellung. Sie steuern, was bei jedem Schritt geschieht – interaktiv, aus dem Code heraus –, sodass Sie schnell iterieren und debuggen können.

Die Aufgabenverteilung sieht wie folgt aus:

Die API ist um drei Clients im SageMaker Python SDKorganisiert, die nacheinander erstellt werden. Ein ServiceClient enthält die Einstellungen auf Plattformebene (Ausführungsrolle, Model-Package-Group, Amazon Simple Storage Service (Amazon S3)-Ausgabespeicherort und optional einen AWS Key Management Service (AWS KMS)-Schlüssel sowie eine Virtual Private Cloud (VPC)-Konfiguration, die Sie auch pro Sitzung überschreiben können) und startet eine Sitzung. Der ServiceClient instanziiert einen TrainingClient (Trainings- und Checkpointing-Operationen wie forward, forward_backward, optim_step, und save_weights_for_sampler), der wiederum einen SamplingClient (sample und einen compute_logprobsinstanziiert (für Generierung und Bewertung während Rollouts und Evaluation). Jeder API-Aufruf entspricht genau einer Operation in Ihrer Schleife. Eine aktive Sitzung stellt zusätzlich einen OpenAI-kompatiblen Chat-Completions-Endpunkt bereit, sodass vorhandener OpenAI-Client-Code daraus sampeln kann.

Abbildung 1: Die Client-Hierarchie des SageMaker Python SDK für Trainings-Sitzungs-APIs

Für bestärkendes Lernenhat die Schleife eine charakteristische Form. Die folgende Abbildung zeigt einen Trainingsschritt und wo jeder Teil ausgeführt wird: Die grünen Schritte rufen APIs auf dienstverwalteten GPUs auf, und die violetten Schritte, Ihre Schleifensteuerung und Belohnungslogik, laufen auf Ihrem Rechner.

Abbildung 2: Ein Trainingsschritt, mit Sampling-, Gradienten- und Checkpoint-Operationen auf verwalteten GPUs sowie Schleifensteuerung und Belohnungslogik auf Ihrem Rechner

Jeder Schritt lädt einen Batch von Prompts und (1) sampelt vom aktuellen Modell auf der GPU. Ihr Python-Code bewertet dann (2) diese Ausgaben: ein Verifier, ein Reward-Model-Aufruf, ein Test-Harness, alles, was Sie in Code ausdrücken können. (3) forward_backward und optim_step die Punktwerte in ein Gradienten-Update auf der GPU umwandeln, und (4) save_weights_and_get_sampling_client() aktualisiert den Sampler, sodass der nächste Schritt aus dem aktualisierten Modell zieht.

Sie besitzen die Daten, die Schleifensteuerung und die Belohnung; SageMaker AI führt die GPU-gebundenen Sampling-, Gradienten- und Checkpoint-Operationen aus. Dasselbe Grundprinzip deckt ab überwachtes Fine-Tuning (SFT) und andere LLM-Trainingstechniken. Für SFT entfällt das Sampling und der Belohnungsschritt, und man ruft forward_backward mit einer cross_entropy Verlust über gelabelte Beispiele hinweg.

Für wen ist das gedacht und wo passt es in Amazon SageMaker AI?

Trainings-APIs sind der richtige Ausgangspunkt, wenn Sie schnelle Iterationszyklen und algorithmische Freiheit benötigen – etwa für einen eigenen Loss, ein neues RL-Reward-Signal oder eine mehrstufige Pipeline – aber nicht selbst GPU-Infrastruktur bereitstellen oder betreiben wollen.

Die Abrechnung erfolgt pro verarbeitetem Token, sodass nur für die Rechenleistung berechnet wird, die Ihre Schleife verbraucht. Wenn Ihre Reward-Funktion 10 Sekunden benötigt, um einen Batch auf Ihrem Rechner zu bewerten, wird keine GPU-Zeit berechnet, während sie läuft. Verglichen mit einer stundenbasierten GPU-Abrechnung bietet dies zwei praktische Vorteile. Erstens entstehen keine Kosten für ungenutzte GPUs: Reward-Berechnung, Datenladen, Protokollierung und Debugging laufen außerhalb der abrechenbaren Zeit, und Ihnen wird nur berechnet, wenn Sie eine GPU-gebundene Operation aufrufen. Zweitens skaliern die Kosten mit Ihrem Datensatz und Ihrem Generierungsvolumen, nicht mit der realen Zeit oder der Clustergröße, was die Ausgaben leichter planbar und Experimente leichter vergleichbar macht. Die Ersparnisse gegenüber der Reservierung von GPUs für die gesamte Sitzung können bei Workloads mit hoher lokaler Rechenlast erheblich sein, etwa bei komplexen Reward-Modellen, Multi-Tool-agentischen Rollouts oder Human-in-the-Loop-Evaluationen.

Trainings-Sitzungs-APIs sind für Teams konzipiert, deren Arbeit durch den Trainingsalgorithmus definiert wird, nicht nur durch den Datensatz. Da Sie die Schleife selbst zusammenstellen, unterstützen dieselben Primitive überwachtes Fine-Tuning (SFT) mit einem cross_entropy Verlust über gelabelten Daten, Bestärkendes Lernen with ppo, importance_sampling, oder cispo (GRPO, PPO und andere gruppenrelative Methoden) und benutzerdefinierte Ziele, wie Präferenzoptimierung oder Destillation, die aus denselben Bausteinen aufgebaut sind. Dazu passen:

• Post-Training- und Alignment-Forscher, die RL-Methoden (GRPO, PPO, Importance-Sampling-Varianten), eigene Advantage-Schätzer oder neue Zielfunktionen implementieren.
• Teams, die Supervised Fine-Tuning oder Preference Optimization durchführen und eine nicht standardmäßige Loss-Funktion, einen Daten-Curriculum oder eine mehrstufige Pipeline benötigen, die konfigurationsgesteuerte Rezepte nicht ausdrücken können.
• Entwickler von agentenbasiertem und mehrstufigem RL, die benutzerdefinierte Rollout-Logik, Reward-Shaping und Curricula benötigen.
• Machine-Learning-(ML)-Engineers, die die konfigurationsgesteuerte Feinabstimmung hinter sich gelassen haben, aber nicht die operative Last von GPU-Clustern übernehmen wollen.
• Forscher, die bestehenden Training-Loop-Code portieren, da das operationsbasierte Programmiermodell direkt auf Loops abgebildet werden kann, die gegen Low-Level geschrieben wurden forward_backward, optim_step, und sample APIs.

„Die unterstützten Basismodelle und die AWS-Regionen, in denen jedes verfügbar ist, sind in der“ Dokumentation der Schulungssitzungs-APIs.

Wo Trainings-Sitzungs-APIs innerhalb von Amazon SageMaker AI einzuordnen sind

Trainings-Sitzungs-APIs nehmen bei den Modell-Anpassungsoptionen in SageMaker AI eine eigene Position ein: deutlich mehr Kontrolle über den Trainingsalgorithmus als vollständig verwaltet Serverless Model Customization, und weit weniger operativer Eigenverantwortung als beim Training mit Amazon SageMaker Training Jobs oder Amazon SageMaker HyperPod.

Abbildung 3: Einordnung der Trainingssitzungs-APIs unter den Modell-Anpassungsoptionen in SageMaker AI

Anatomie einer Trainingsschleife: GRPO bei Matheaufgaben

Dieser Abschnitt zerlegt eine minimale Reinforcement-Learning-Schleife, um zu zeigen, wie die Operationen zusammenwirken. Die Aufgabe ist GSM8K, ein Benchmark aus etwa 8,500 Mathe-Textaufgaben der Grundschule, jeweils mit einer einzigen numerischen Antwort, wie im folgenden Beispiel:

Natalia verkaufte im April Clips an 48 ihrer Freunde und im Mai verkaufte sie halb so viele Clips. Wie viele Clips verkaufte Natalia insgesamt im April und Mai? → 72

Wir trainieren einen LoRA-Adapter für OpenAIs GPT-OSS 20B mit Group Relative Policy Optimization (GRPO): für jede Frage erzeugt das Modell mehrere Antworten, jede Antwort wird auf Korrektheit bewertet, und Antworten, die den Gruppendurchschnitt übertreffen, werden verstärkt. Das vollständige, lauffähige Beispiel befindet sich in der Amazon SageMaker Examples-Repository.

Bevor Sie eine Sitzung öffnen, führen Sie die Voraussetzungen in der Dokumentation durch.

Sitzung öffnen

Erstellen Sie eine Sitzung, holen Sie sich dann den Tokenizer und ein SamplingClient daraus:

from sagemaker.train.training_session import ServiceClient, SamplingParams, AdamParams

svc = ServiceClient(role_arn=role, region=REGION)

training_client = svc.create_lora_training_client(
    "openai/gpt-oss-20b",                        # a name from svc.list_base_models()
    training_session_name=SESSION_NAME,
    rank=32, alpha=64,                           # LoRA rank and scaling factor
    model_package_group_arn=MODEL_PACKAGE_GROUP_ARN,
    s3_output_path=OUTPUT_S3_PATH,
)
tokenizer = training_client.get_tokenizer()
sampling_client = training_client.create_sampling_client()

Definieren der Belohnung und des Vorteils

Die aufgabenspezifische Logik besteht aus zwei kurzen Funktionen, die auf Ihrem Rechner ausgeführt werden:

def compute_reward(response_text, ground_truth):
    predicted = extract_answer(response_text)      # parse the number after "#### "
    return 1.0 if predicted == ground_truth else 0.0

def compute_advantages(rewards):
    rewards = np.array(rewards)
    advantages = rewards - rewards.mean()          # center within the group (the GRPO baseline)
    return advantages / (rewards.std() + 1e-8)     # normalize

Die Schleife

Jeder Schritt folgt den vier in Abbildung 2 dargestellten Phasen. draw, build_prompt, und build_training_datum sind kurze Hilfsfunktionen im Beispiel:

for step in range(NUM_STEPS):
    datums, pending = [], []
    # (1) sample K responses per question in one request each (GPU); returns a future right away
    for question, ground_truth in draw(PROMPTS_PER_STEP):          # data loading is local
        prompt = tokenizer.encode(build_prompt(question), add_special_tokens=False)
        pending.append((prompt, ground_truth, sampling_client.sample(
            prompt, num_samples=K,
            sampling_params=SamplingParams(max_tokens=512, temperature=0.7, top_p=0.9))))

    for prompt, ground_truth, future in pending:
        result = future.result()                                    # wait once the batch is in flight
        # (2) score each response, then center the rewards within the group (local)
        rewards = [compute_reward(tokenizer.decode(s.tokens), ground_truth) for s in result.sequences]
        advantages = compute_advantages(rewards)
        datums += [build_training_datum(prompt, s, a) for s, a in zip(result.sequences, advantages)]

    # (3) one gradient step over the batch (GPU); the optimizer step is sequenced behind it
    fwd_bwd_op = training_client.forward_backward(datums, loss_fn="ppo")
    optim_op = training_client.optim_step(AdamParams(learning_rate=1e-5))
    fwd_bwd_op.result(), optim_op.result()

    # (4) refresh the sampler so the next step's rollouts come from the updated policy (GPU)
    sampling_client = training_client.save_weights_and_get_sampling_client()

training_client.stop()

Drei Details machen diese Schleife funktionsfähig:

Operationen geben Futures zurück. Jeder GPU-gebundene Aufruf gibt sofort ein APIFuture zurück: Die Anfrage wird in eine Warteschlange gestellt und von einem Hintergrund-Submitter gesendet, und .result() wartet auf den Wert. Die Schleife sendet alle Sampling-Anfragen eines Schritts, bevor sie auf eine davon wartet, sodass der gesamte Batch gleichzeitig generiert wird. Trainingsoperationen werden in der Reihenfolge ausgeführt, in der Sie sie aufrufen, sodass optim_step abgesendet werden kann, bevor forward_backward abgeschlossen ist.

Das Datum enthält alles, was der Loss benötigt. build_training_datum packt die Prompt- und Response-Tokens, die per-token logprobs von sample(), den Advantage und eine weights Maske, die den Prompt aus dem Loss heraushält.

Checkpoints sind Modellpakete. Jedes Speichern registriert eine Version in Ihrer Modellpaketgruppe, einschließlich der Per-Step-Sampler-Aktualisierung in der obigen Schleife. Nach dem Training fügen zwei weitere Speicherungen ein zusammengeführtes Modell für die Bereitstellung und einen fortsetzbaren Zustand mit Gewichten und Optimizer hinzu.

training_client.save_weights_for_sampler(merge_weights=True).result()   # merged, deployable model
training_client.save_state().result()                                   # resumable state (weights + optimizer)

Um fortzufahren, starten Sie eine neue Sitzung vom letzten Checkpoint mit create_training_client_from_state_with_optimizer.

Eine laufende Sitzung wird auf Ihr Kontingent für gleichzeitige Sitzungen angerechnet, bis Sie stop() aufrufen oder sie abläuft; beenden Sie sie daher, wenn Sie fertig sind, auch in Fehlerpfaden, oder öffnen Sie den Client in einem with Block. Der zusammengeführte Checkpoint wird für einen Real-Time-Endpunkt bereitgestellt. Um Checkpoints stattdessen als LoRA-Adapter bereitzustellen, siehe das Cookbook-Kapitel deployment notebook.

In unserem Lauf des Beispiels steigerten 10 Trainingsschritte auf 40 Problemen die Held-out-Genauigkeit von 56 Prozent auf 75 Prozent.

Ein realer Arbeitslast-Fall: RL auf Competitive Programming

GSM8K hält die Mechanik sichtbar, aber eine produktive RL-Arbeitslast kommt mit mehr daher: längere Generierungen, Interaktion über mehrere Runden, einen externen Grader und genügend Durchsatz, um die GPUs auszulasten. Das SageMaker training sessions cookbook verpackt diese Technik in wiederverwendbare Trainingsschleifen auf Basis derselben Operationen. Für eine neue RL-Aufgabe schreibt man eine Umgebung: was das Modell sieht, wie seine Antwort bewertet wird und ob es eine weitere Runde bekommt. Das Cookbook übernimmt Batching, Chat-Templates, Advantage-Schätzung, überlappendes Sampling und Training, Validierung, Checkpointing mit Wiederaufnahme sowie Metriken pro Schritt und zeichnet die Konfiguration jedes Laufs auf, damit er wiederholt werden kann. Das code_rl -Rezept nutzt es, um ein Modell auf Competitive-Programming-Problemen zu trainieren.

Die Aufgabe. Die Probleme stammen aus dem DeepCoder-Datensatz (PrimeIntellect, TACO und LiveCodeBench für das Training; Codeforces und LiveCodeBench für das Testen). Das Modell schreibt eine Python-Lösung, die gegen die Unit-Tests des Problems in einem SandboxFusion -Container auf Ihrer Maschine ausgeführt wird. Die Belohnung ist 1, wenn alle Tests bestehen, und sonst 0, mit einem Abzug von 0.1, wenn die Antwort keinen Codeblock enthält.

Multi-Turn-Tool-Nutzung. Mit zwei Runden kann das Modell ein check_solution -Tool aufrufen, die Testergebnisse lesen und seine Antwort überarbeiten, bevor die Endnote erteilt wird. Die Tool-Aufrufe und -Ergebnisse werden im nativen Chat-Format des Modells gerendert, sodass das Modell auf demselben Token-Format trainiert wird, das es bei der Inferenz verwendet.

Was das Rezept zur minimalen Schleife hinzufügt:

• Überlappendes Sampling und Training: Die Rollout-Worker generieren weiter, während der Trainer einen Schritt ausführt. Eine Gruppe von Antworten, die von einer älteren Politik gesampelt wurde, wird verwendet, wenn sie höchstens k Schritte alt ist, andernfalls erneut in die Warteschlange gestellt, und die cispo -Loss korrigiert für den Politikunterschied.
• Streamed Gradientenberechnung: forward_backward wird in Sub-Batches eingereicht, sobald Gruppen fertig sind, statt einmal pro Schritt.
• Held-out-Validierung, Checkpoint-Wiederaufnahme und Metriken pro Schritt, optional veröffentlicht in SageMaker managed MLflow.

Abbildung 4: Das Code-RL-Setup, mit der lokalen Trainingsschleife und Sandbox links und den verwalteten GPUs rechts

Wir trainierten einen LoRA-Adapter (rank 32) für GPT-OSS 20B über 50 Schritte, jeweils mit 128 Problemen und 8 Antworten pro Problem, mit zwei Runden und einer Staleness-Grenze von k=3. Die Trainingsbelohnung stieg von 0.46 auf 0.60, und die Genauigkeit auf 200 held-out-Testproblemen stieg von 40 Prozent auf 50.5 Prozent.

Um es auszuführen, starten Sie die Sandbox und starten das Rezept:

bash tools/sandbox_up.sh
python -m training_sessions_cookbook.recipes.code_rl.train \
    --config training_sessions_cookbook/recipes/code_rl/configs/coderl_oss20b_service.json \
    --max_steps_off_policy 3 --stale_policy requeue --num_rollout_workers 80 --max_steps 50 \
    --base_model_arn "$BASE_MODEL_ARN" \
    --model_package_group_arn "$MODEL_PACKAGE_GROUP_ARN" \
    --role_arn "$ROLE_ARN" \
    --output_s3_path "$OUTPUT_S3_PATH"

Das Kochbuch enthält weitere Rezepte mit derselben Struktur, darunter überwachtes Fine-Tuning auf Chat-Datensätzen (Tulu 3, UltraChat), RL auf GSM8K in vollem Umfang und KTO-Präferenz-Ausrichtung aus binärem Feedback. Jedes Rezept-README listet seine Befehle, Hyperparameter und erwarteten Metriken auf.

Fazit

Die SageMaker AI Training-Session-APIs verlagern die Abstraktionsgrenze hinunter auf die Trainingsoperation: Sie schreiben die Trainingsschleife in Python auf Ihrem eigenen Rechner, und der Dienst führt sie auf verwalteten GPUs aus und übernimmt das Laden von Modellen und LoRA, die verteilte Ausführung, die Fehlerbehebung und das Checkpointing. Eine vollständige GRPO-Schleife reduziert sich auf vier Operationen, sample, forward_backward, optim_stepund save_weights_and_get_sampling_client, wobei sich die Reward- und Advantage-Logik in Ihrem eigenen Code befindet. Dieselben Operationen skalieren auf Multi-Turn-RL mit Tool-Nutzung und Sandboxed-Bewertung, wie das Code-RL-Rezept zeigt.

Eigene Rewards, agentische Rollouts, schrittweises Debugging und neue veröffentlichte Techniken verwenden alle dasselbe Programmiermodell: Der Algorithmus bleibt in Ihrer Schleife, und SageMaker AI betreibt die GPUs.

So starten Sie:

• Wählen Sie ein Rezept aus dem SageMaker training sessions cookbook, zum Beispiel Code RL.
• Führen Sie das GSM8K-Beispiel aus dem Amazon-SageMaker-Beispiel-Repository aus und ersetzen Sie die Reward-Funktion durch Ihre eigene.
• Lesen Sie die Dokumentation der Training-Session-APIs zu Voraussetzungen, Kontingenten, Regionen, unterstützten Modellen und der API-Referenz.

Danksagungen

Dank gilt Arun Nagarajan, Xichao Wang und Shruti Sharma für ihre Beiträge und ihre Anleitung zu diesem Beitrag.

Über die Autoren