Share training code between backends

[Kovbo]

We have a lot of similar training code implemented differently across ART and the serverless backend. That makes RL/SFT changes hard to reason about and easy to drift.

This PR moves the reusable training/runtime logic into ART and makes the serverless backend import it instead of reimplementing it.

Before

We had duplicated training logic in multiple places:

• ART local Megatron had its own runtime setup and RL worker loop
• Serverless Megatron had its own runtime setup plus separate RL and SFT loops
• Serverless Unsloth had its own RL/SFT execution logic
• ART local/backend and serverless/backend each built RL config objects and aggregated training metrics separately

So changing training behavior often meant touching multiple implementations.

How It Works Now

Shared backend glue

Shared RL config construction and metric aggregation now live in:

• src/art/_backend_training.py

Both of these now use that shared helper:

• src/art/local/backend.py
• src/art/serverless/backend.py

Shared Megatron job protocol and execution

Megatron’s shared cross-repo pieces now live in ART:

• src/art/megatron/train.py
• src/art/megatron/client.py
• src/art/megatron/jobs.py
• src/art/megatron/sft_batches.py
• src/art/megatron/merge.py
• src/art/megatron/runtime_env.py

The public worker API is now in src/art/megatron/train.py, not a separate shared.py.

That module owns the actual Megatron worker execution:

• runtime/model/optimizer setup
• RL loop
• SFT loop
• LoRA + optimizer load/save
• metrics logging
• job finalization and cleanup

The orchestration layers are now thin wrappers around that API.

Shared Unsloth execution

Shared Unsloth execution now lives in:

• src/art/unsloth/train.py

That module exposes the reusable serverless-facing API:

• create_unsloth_train_context
• run_unsloth_rl_training
• run_unsloth_sft_training

Serverless now imports those directly instead of maintaining a separate implementation.

Design Choice

Megatron orchestration is still separate from Megatron execution.

The shared ART execution layer knows how to run Megatron jobs, but process lifecycle / handoff / sleep-wake behavior stays in the
callers:

• ART local orchestration lives in src/art/megatron/service.py
• Serverless Megatron orchestration lives in trainers/megatron_trainer.py and megatron/train.py

This keeps the shared code focused on training logic rather than backend-specific worker management.

Current End-to-End Flows

ART local RL with Megatron

1. LocalBackend.train() / _train_model() build config and metrics via src/art/_backend_training.py
2. MegatronService.train() pauses vLLM, ensures the Megatron process is running, and writes a job file via src/art/megatron/ client.py
3. The Megatron worker loop in src/art/megatron/train.py picks up that job and runs run_megatron_rl_job()
4. MegatronService publishes the resulting checkpoint and re-registers the LoRA for inference

ART local SFT with Megatron

1. LocalBackend._train_sft() tokenizes trajectories into SFTBatch
2. MegatronService.train_sft() materializes those batches to disk and writes a MegatronSFTTrainingJob
3. The same worker loop in src/art/megatron/train.py runs run_megatron_sft_job()

Serverless RL

1. The workflow restores model state/artifacts and chooses a trainer via trainers/__init__.py
2. MegatronTrainer.train() or UnslothTrainer.train() delegates into ART shared execution
3. Megatron uses ART’s shared job schema/client and ART’s worker loop via megatron/train.py
4. Unsloth uses ART’s shared train context and RL helper directly

Serverless SFT

1. The workflow tokenizes uploaded training data into SFTBatch objects
2. MegatronTrainer.train_sft() uses ART’s shared materialize_sft_batches() and MegatronSFTTrainingJob
3. The Megatron worker loop runs run_megatron_sft_job()
4. UnslothTrainer.train_sft() calls run_unsloth_sft_training() from ART directly

Operational Note

For Megatron in serverless, local /tmp/megatron_training_jobs files are only a per-worker handoff between the trainer and the Megatron
subprocess. They are not the customer-facing queue.

Customer backlog still lives in Temporal.

This PR also clears stale local Megatron job JSONs before enqueueing a new job in serverless so an interrupted worker does not
accidentally replay an old local job file.

Result

After this change:

• training execution logic lives in ART
• serverless imports ART’s Megatron and Unsloth training code instead of reimplementing it
• RL backend config/metric aggregation is shared
• Megatron RL and SFT use one shared worker implementation
• changes to training behavior mostly happen in one place

[FurtherAI]

I will propose a merge with main after #619. We'll need to collect the changes I made to train.py mostly and settle on a unified API.

The rest of it looks good! Though I'll test the actual megatron training and ensure the merge still keeps correctness and performance. You should extend the correctness tests to SFT though.

[FurtherAI]

1. The PR now signals "all done" before offloading Megatron state. The sentinel is written in shared.py, while offload happens later in train.py. On main, offload happened first at train.py and only then did rank 0
   emit completion at train.py. That can reintroduce same-node wake/OOM races.
2. Missing adapter checkpoints no longer fail hard. The new shared loader in shared.py
   silently resets LoRA params instead of raising. main hard-fails on a missing adapter.
3. The log_path API cutover is still partial. The shared job schema exposes log_path in jobs.py, and shared execution uses it in shared.py, but the local controller still tails
   and deletes the default path in service.py.

All the good points found by codex. I'll add to 1 that the some objects such as the packed tensors need to be deleted before the dirs can be removed, since they are memory mapped and thus hold the files open. So all object deletion should happen before dir removal.

Correctness tests were passing, which is good.

[Kovbo]

Good catch.

I think everything has been updated now, and I also added some additional training loop refactoring.

My Codex also flagged a few other issues, but they do not appear to be regressions in this branch. I am not sure whether they are real issues, though:

1. /art/megatron/train.py:355 still requires grad_accumulation_sequences % dp_world_size == 0, and the public config builders still leave it at the default 1 in /art/_backend_training.py:15. So DP>1 Megatron can break through backend issue, just not caused by our sharing refactor.
2. In /art/megatron/train.py:495-540, reduced_loss is computed as sum of per-rank averages, not a global token-weighted average. That makes DP>1 logging wrong.

[FurtherAI]

1. /art/megatron/train.py:355 still requires grad_accumulation_sequences % dp_world_size == 0, and the public config builders still leave it at the default 1 in /art/_backend_training.py:15. So DP>1 Megatron can break through backend issue, just not caused by our sharing refactor.
2. In /art/megatron/train.py:495-540, reduced_loss is computed as sum of per-rank averages, not a global token-weighted average. That makes DP>1 logging wrong.

1. Yeah, let's set the default grad_accumulation_sequences to None and if none, update it to dp_world_size here. The assert is intentional though, so it stays.
2. As for this, finalize_model_grads in finalize_model_grads_extended correctly reduces this across ranks I believe. It would be good to drop a comment there though ("# num_tokens is reduced in place across ranks").

The previous issues seem to be solved. Looks good

Did you add SFT to the correctness tests? Or is that why supports_sft=False?

[FurtherAI]

Changes look good, I ran the correctness tests and they are passing.
Last thing is the plan for SFT correctness tests?

[FurtherAI]

I modified the api a bit to be able to run the correctness tests. Big changes are:

• has a separated out training step, which matches rl pretty closely
• matches rl gradient accumulation so data parallelism works correctly

Small fix:

• optimizer states keyed by objective (SFT vs RL) so back to back jobs wouldn't overwrite the other's optimizer.

The correctness/sensitivity tests are passing so looks pretty good. Could you check that it runs end-to-end? (trainability will be broken, bug is in offload, but I'll push a fix for that in my next PR).