RFC: Checkpointing Beyond Model Weights (Why we can’t do it now, and what we’ll do next)

[DNXie]

TLDR: We can’t (cleanly) checkpoint everything yet (e.g., dataloader, replay buffer, RNG, etc.) because Titan’s checkpointer lives inside trainer.engine and exclusively controls the step-<N> folders, while those other pieces live in separate actors. There’s no safe, centralized way to co-write into the same step right now, so we will need to deferring full multi-component checkpointing until after PTC.

What does work today: with #444, we can enable saving and resuming model weights + optimizer + LR scheduler via Titan’s checkpointer.

Scope: This RFC only discuss the challenges in saving the checkpoint. For loading, the discussion can be referred to #425

---

1) Context (today’s flow)

We spin up components in main:

(
    dataloader,     
    policy,       
    trainer,        # has ForgeEngine
    replay_buffer,  
    compute_advantages,  
    ref_model,       
    reward_actor,    
) = await asyncio.gather(...)

Model checkpointing is delegated to TorchTitan via the trainer’s engine:

• trainer creates the engine and loads a checkpoint:

self.engine = ForgeEngine(ForgeJobConfig(**engine_config))
self.engine.checkpointer.load(step=self.step)

• Every training step (in train_step), the trainer asks Titan to save:

self.engine.checkpointer.save(
    curr_step=self.step,
    last_step=self.step == self.num_training_steps,
)

Titan’s checkpointer writes model weights, optimizer state, LR schedulers into:

<folder>/step-<N>/__0_0.distcp

For example

./checkpoint/step-100/__0_0.distcp

Per issue #362, we also want to save/load the states for:

• data step,
• replay buffer data,
• RNG states,
• etc.

---

2) The problems

Problem 1: Step-folder ownership

We have Titan-owned directory per saving step (e.g., step-200) created from inside trainer.engine.checkpointer. Other actors (dataloader, replay buffer) do not have access to the trainer’s engine or to Titan’s private folder-naming method. That leaves us with two awkward choices:

1. Two folders per step

   checkpoint/
     step-200/          # Titan
       __0_0.distcp
     step-200-other/    # Ours
       dataloader.json
       replay_buffer.bin
       rng.json
   

   Downsides: clunky UX, hard to purge/retain atomically, and easy to drift.

2. Single folder per step (preferred)
   To co-locate our files inside step-200/, we must either:

   • call Titan’s private _create_checkpoint_id to learn the folder name, but other components (e.g., dataloader) doesn't have an engine.
   • re-implement a look-alike function and hope it never diverges.
   • (preferred) Add a path parameter to the checkpointer.save.

Problem 2: No unified checkpoint scope

Currently, non-model states (like dataloader, replay_buffer, RNG, etc.) live in separate actors/services (e.g., dataloder), and the trainer, which owns Titan’s checkpointer, only manages model, optimizer, and LR scheduler.

Because Titan’s checkpointing is embedded inside trainer.engine.checkpointer, there’s no single coordinator that can write all components into the same step-<N> folder in a clean, synchronized way.

---

3) Proposal for Problem 2

Option 1: Make trainer own all other components

class RLTrainer:
    self.dataloader = ...
    self.replay_buffer = ...
    self.rng = ...

This is very fast to implement. However, it introduces heavy coupling, breaks actor/service boundaries, hurts scalability and reuse.

Option 2: Reimplement checkpointing ourselves

Write our own model/optim/lr checkpointing.

It gives us full control over layout and atomicity.

Downside

• Rebuilding the hardest part (distributed model/optim/lr, DCP, async, FT).
• High risk, high effort; guaranteed divergence from Titan.

Option 3: Coordinator layered above current actors

Introduce a light Checkpoint Coordinator that:

• calls Titan to save the model/lr/optimizers to specified path (requires one small API addition to Titan save: path=),
• then asks each actor to state_dict() and writes to the same folder (e.g., step-200/dataloader.json, etc.),
• on load, after Titan resolves the step it will load, coordinator tries to load each components' states by calling their load_state().

class CheckpointCoordinator:
    def __init__(self):
        self._trainer: RLTrainer = None        
        self._components: Dict[str, ForgeActor | SeverceInterface] = {}        

    def set_trainer(self, trainer:RLTrainer):
        self._trainer = trainer

    def register(self, name, comp: ForgeActor | SeverceInterface):
        self._components[name] = comp

    async def save(self, step: int, path: str):
        path = get_path(folder, step)
        if self._trainer:
            self._trainer.engine.checkpointer.save(path = path)  
        for name, comp in self._components.items():
            states = comp.state_dict()
            dump_json(states, f"{path}/{name}.json")

    async def load(self, step: int, path: str):
        ...

The changes required in grpo/main:

coord = CheckpointCoordinator()
coord.set_trainer(trainer) 

coord.register("dataloader", dataloader)
coord.register("replay_buffer", replay_buffer)
...

await coord.load(step, path=checkpoint_folder)

async def continuous_training():
     ...
     await coord.save(step, path=checkpoint_folder)

It is a relatively easy change and keeps most of the existing wheels.

Downside

• Nested structure: we still do coord.save which calls self._trainer.engine.checkpointer.save.
• Slightly specific to our the existing grpo script. May have generalizability issue in the future.

Option 4: Standalone ForgeCheckpointManager

Longer-term, we could create a standalone manager ForgeCheckpointManager that inherits Titan’s CheckpointManager and orchestrates both Titan and forge's components in one save()/load() call. Actors register their export_state/import_state with this one object; main calls just it.

Open question: where does ForgeCheckpointManager live if the engine stays inside the trainer? And how does it read/write model/optim/lr state without re-nesting the trainer or breaking actor decoupling?

[ebsmothers]

Thanks for putting this together! This RFC is awesome and very thorough. Regarding problem 1 and the preferred approach

Add a path parameter to the checkpointer.save

@tianyu-l @wwwjn is there anything preventing us from doing this today on the titan side?

Problem 2 is tricky. My mental model is that checkpointing is kind of analogous to what we do with metric logging: we ultimately need some sort of global singleton that can delegate to subcomponents for different workers, but is able to coordinate and manage their state. I know in our discussion yesterday I was kinda pushing for option 4, but looking at the code I actually like option 3. I'd definitely like to see this path further fleshed out. E.g. imo it'd be preferable if we could have a bit more flexibility on saving different components at different cadences (or from within the actors themselves). Personally I think we shouldn't go for either option 1 or 2, I think the cons you called out are enough to preclude serious consideration for either approach.

Would be very interested to hear some other folks' thoughts as well. Tagging @joecummings @JenniferWang @pbontrager @allenwang28 since I suspect you all will have opinions here too.

[tianyu-l]

call Titan’s private _create_checkpoint_id to learn the folder name, but other components (e.g., dataloader) doesn't have an engine.

Sorry is this doable or not? If it's doable I think it's fine.

re-implement a look-alike function and hope it never diverges.

No let's not do this.

(preferred) Add a path parameter to the checkpointer.save.

I'm OK with adding a kwarg (default to None), but I feel users shouldn't be incentivized to be creative on such implementation details. So I'd prefer option 1.

[DNXie]

@tianyu-l

call Titan’s private _create_checkpoint_id to learn the folder name, but other components (e.g., dataloader) doesn't have an engine.
Sorry is this doable or not? If it's doable I think it's fine.

This is technically doable but would be super not ideal to do. Because each components (e.g., trainer, dataloader, replay_buffer) are supposed to be independent/isolated.

No let's not do this.

Agree.

I'm OK with adding a kwarg (default to None), but I feel users shouldn't be incentivized to be creative on such implementation details. So I'd prefer option 1.

Thanks for the flexibility! Regarding option1, same to my points above, it would introduce heavy coupling and hurt the scalability and reuse. It is not an ideal option for the overall design.

[casteryh]

standalone manager ForgeCheckpointManager that inherits Titan’s CheckpointManager and orchestrates both Titan and forge's components in one save()/load() call.

Not that I am against a standalone ForgeCheckpointManager, but I think "inheriting" Titan's is a bad idea.

[casteryh]

Problem 1: Step-folder ownership

To play devil's advocate. The following structure is actually acceptable to me:

/checkpoints
|-- titan
|   |-- step-1
|   |-- step-2
|-- component_1
|   |-- step-1
|   |-- step-2

as long as we have control on which step each component is saving to / loading from.

In particular, different components's last checkpoints might be at different steps because of potential failures etc.

This could work if titan is to make the CheckpointManager::_find_load_step api public/stable.

[tianyu-l]

@DNXie oh sorry I misunderstood the problem earlier.

(preferred) Add a path parameter to the checkpointer.save.

Yes, I think if you are calling multiple checkpointer from forge, we should clearly do this.

But what's not clear to me is how you handle the checkpointing of other components outside titan engine. Could you share some pointer?

Fwiw, for RNG states we should checkpoint the trainer's RNG states in torchtitan checkpoint manager, although this is not supported today.

[DNXie]

@tianyu-l

But what's not clear to me is how you handle the checkpointing of other components outside titan engine. Could you share some pointer?

We don't have this part implemented yet. This RFC is just to discuss the plan. If you see 3) Proposal for Problem 2, I listed four options. And @ebsmothers prefers option3. I agree with him that we shouldn't go for option 1 or 2.

Fwiw, for RNG states we should checkpoint the trainer's RNG states in torchtitan checkpoint manager, although this is not supported today.

That's wonderful. Although other components (e.g., dataloader) may have their own randomness.

[tianyu-l]

@DNXie
I see. My next question is

Currently, non-model states (like dataloader, replay_buffer, RNG, etc.) live in separate actors/services (e.g., dataloder),

For these components living in separate actors/services, do they have distributed states or create DTensor?

• If so, are the distributed degrees the same as the trainer?
  • If so: I don't see why they're not part of the trainer and why we don't do option 1.
  • If not: then what degrees do they have? Are you going to call dcp.save/load by yourselves? I think that's going to be error-prone and the semantics are not clear (with / without resharding).
• If not, option 3 sounds the right thing to do, and option 4 sounds overkill.

[allenwang28]

Thinking out loud, from #362 here are the things that we should be checkpointing:

• Data loader (the index of the data loader, NOT the same thing as the Titan data loader)
• Model weights, optimizer, LR scheduler (completely owned by Titan's data loader)
• Replay buffer (data stored in the replay buffer)

It's a mistake to have all of it be owned by titan, although it probably makes sense semantically to use titan's train step as the synchronization point for global checkpointing.

So in that case I would:

1. Leverage Titan's checkpointing as is for model weights, optimizer, LR scheduler
2. Checkpoint our data loader in the same file

which leaves replay buffer checkpointing. Our replay buffer likely needs a design re-visit because it's still very simple - we might want to actually store data in disk or a database, because its lifecycle should not be ephemeral with the job if we want to e.g. examine the state of the replay buffer afterwards.

I mention all this to say - if we start with checkpointing trainer state with titan and checkpointing the data loader, that is probably a good enough starting point

[DNXie]

@allenwang28

Checkpoint our data loader in the same file

If so, we need to save data loader states in titan's checkpointer manager. Are you suggesting we should let trainer be aware of the dataloder? Like:

class trainer(ForgeActor):
     dataloader: DatasetActor

[allenwang28]

nope, more something like

class DataLoader(Actor):
    def checkpoint(self, folder: str, step: int):
        ...