Skip to main content
Use algorithm = "custom" for DPO, distillation, or another objective that needs its own training loop. Your Python script runs on your machine and calls a managed Flash GPU for sampling, scoring, backward passes, and optimizer steps. You own data loading, prompt formatting, rewards, and logging. Custom training currently supports text-only Qwen/Qwen3.5-9B on one Modal H100, with up to 4,096 tokens per sequence and LoRA rank 64. Other models, images, multi-GPU shapes, and the built-in thinking switch are not supported in this mode. Custom distillation scripts supply their own teacher access.

​
Start from a DPO example

Install the custom-loss extra in the Python environment running the CLI, then create an editable script and matching config:
flash custom init creates algorithm.py and algorithm.toml, preserving existing files by refusing to overwrite them. The starter demonstrates paired chosen/rejected answers, a frozen reference distribution, a local PyTorch DPO loss, metrics, and checkpoint state. Replace its toy pairs with your data.
algorithm.toml
The script path is relative to the config file. Do not add [environment], [wandb], or thinking: the local script owns those concerns. Set learning rates in client.optim_step(...), not in [train]. See Custom configuration for the accepted fields.
Both previews require authentication and a server advertising custom-training support. They validate without allocating a GPU. The quote is a session cost cap based on max_wall_seconds, not a prediction of your script’s runtime.

​
Write the training loop

Define async def main(client) for nonblocking local control, or def main(client) for blocking calls. For the async form, await every operation: Construct batches with from flash.custom import Datum. A Datum holds a list of token IDs and optional weights. Position i scores the next token, tokens[i + 1], so each row has len(tokens) - 1 log probabilities or weights. A batch contains at most 16 rows. Each row has at least two tokens and must fit train.max_context_tokens. forward_backward requires explicit weights and computes the sum of -weight * logprob across tokens and rows. Set prompt-target weights to zero and apply your own averaging denominator. Gradients accumulate until optim_step or zero_grad. For forward_backward_custom, define loss_fn(data, logprobs) returning a finite scalar PyTorch tensor and a dictionary of finite numeric metrics. The provided log-probability tensors support local autograd, including losses that compare multiple sequences. The generated DPO starter shows this pattern. Operations within one session execute in order. Async control lets other local work run while awaiting the GPU; it does not run simultaneous optimizer mutations. If a request fails with an unresolved outcome, call retry_pending() before issuing another operation. This retries the pending operation without applying it twice.

​
Keep the controller running

The flash train process is your controller. Run it in a persistent terminal such as tmux; --background is rejected for custom training. Ctrl-C cancels the run and preserves completed checkpoints. An unhandled script exception also cancels it. This differs from SFT and GRPO, where Ctrl-C only detaches. Flash reserves the GPU while your script prepares data or waits for external rewards. Charges use measured session time, capped at the accepted wall-time quote. The session also stops at its idle or wall-time limit. See Custom-session billing. Use the normal run commands from another terminal:

​
Resume a checkpoint

Pass a JSON-serializable algorithm_state object of at most 1 MiB to optim_step so periodic saves include your data cursor, reference scores, and other algorithm state. Automatic saves occur on the save_every cadence, at save_at_steps, and at the final step; the final adapter is also published. The script is responsible for restoring its own state from client.restored_state and continuing from client.step. To continue, submit a new run with the exact saved step:
Use a custom checkpoint from the same base model, with matching LoRA rank and alpha. The new max_steps must exceed the restored step and is the total target, not the number of additional steps. /final and train.init_from_adapter are not custom-resume selectors. In-place load_checkpoint is unavailable for managed sessions; use custom.resume_from for durable continuation. Custom runs do not automatically retry a failed local controller. Restart from a completed checkpoint. Deploy or download the saved adapter with the normal checkpoint workflow.