Publishing weights in to torchstore from RLTrainer and getting them from policy engine.

[pradeepfn]

Working integration test that works with Llama8B.

1. TorchTitan publishes weights to torchstore (TS),
2. vLLM service consumes weights from TS.

State-dict conversion to HF format happens during weight-save/publish in to TS using sd_adaptor (mostly). Custome param concatanations that are not supported by the current sd_adaptor impl of the model was implemented directly within the weights publish function.

[LucasLLC]

@pradeepfn please clarify, does put hang or throw?

[pradeepfn]

@pradeepfn please clarify, does put hang or throw?

No torchstore issues.

[pbontrager]

This looks really great! Glad we can finally have this started. Mostly left questions and a few nits.

[pbontrager]

Why are we starting at 1? Also, we probably want a todo to update this from the checkpoint

[pradeepfn]

Because policy engine starting at 1. Lets keep this fragile contract as it is. The true version has to come from a config or external book-keeping entity.

[pbontrager]

I don't think we can change this without risking breaking checkpoint expectations from titan side. I'd rather just use a separate variable in the trainer for "checkpoint name" (can be a property that's just current_step + 1 for now). This could also be passed in from the controller which would be better.

[pbontrager]

Where is this coming from? When you call this, does it create the sd right then or did it have to be saved in the train step earlier? Does it return the sd on GPU or CPU? Also does it handle blocking the trainer from updating the weights while it's getting them?

[pradeepfn]

I'm accessing the module state-dict prepped by torch.titan as part of checkpoint save.

1. This is a in-memory state-dict. ( Tensor/DTensor).
2. It returns tensors with original storage. Means GPU/UVM backed tensors.

Also does it handle blocking the trainer from updating the weights while it's getting them?

Hmm.. it does not block the trainer. However, ForgeEngine drive the trainer using train_step. Therefore, there is no race-conditions with current code.

There is improvements to be made to this code. In the ideal case;

1. the state-dict get prepped for weight-exchange and checkpoint save purposes.
2. Once the initial state-dict prep we can cache the prepped state-dict for later iterations of the training steps for efficiency reasons ( if there is opportunity).
3. We move all the model weights and optimizer state to torchstore.
4. Policy engine (only) lookup the model-weights from torchstore
5. Async checkpointing upload lookups model-weights and optimizer states for uploading in to remote persistent storage.

We don't have all the piece right now. But tapping in to checkpoint state-dict is the right thing to do as the first step.

[pbontrager]

I guess you're right that it should be mostly safe since we control the update from the controller. But since they're async calls they could be overlapped so we'll have to be careful for now.

[allenwang28]

awesome!

[allenwang28]

When would "model" not be in the self.engine.checkpointer.states? In other words, can we update the assertion error to be more informative? Does this fail if the user didn't initialize the trainer properly/what do they need to do to make it work??

[pradeepfn]

Essentially, this only happens if the checkpoint_manager of torchtitan is not initialized prior to calling push_weights routine. I can update the error message (followup PR) with that.

[allenwang28]

is this new rtol/atol expected? question for @pbontrager

[pbontrager]

This is never great, but given the bf16/fp16 comments I could see that. This is also an allclose and not a comparison of the mean so we should be safe here. If we can load the hf side with bf16 instead of fp16 we might be able to regain the tighter tolerance.

[allenwang28]

maybe we should hardcode the trainer config here rather than load from apps/rl/llama3_8b.yaml

[pbontrager]

Approving this with a few nits

[pbontrager]

I don't think we can change this without risking breaking checkpoint expectations from titan side. I'd rather just use a separate variable in the trainer for "checkpoint name" (can be a property that's just current_step + 1 for now). This could also be passed in from the controller which would be better.

[pbontrager]

I guess you're right that it should be mostly safe since we control the update from the controller. But since they're async calls they could be overlapped so we'll have to be careful for now.

[pbontrager]

This is never great, but given the bf16/fp16 comments I could see that. This is also an allclose and not a comparison of the mean so we should be safe here. If we can load the hf side with bf16 instead of fp16 we might be able to regain the tighter tolerance.

[pradeepfn]

FYI: @joecummings is working on merging this diff. This is because of recent API changes that was pushed by him. @joecummings can add more insights. thanks.

[Ritesh1905]

Question: are there benfits to doing this at the state dict level and not at the key level where we could parallelize the individual put operation per key?

[joecummings]

Right now, the benefit is simplicity.

Eventually I imagine we will want to do this on a per-slice level.

cc @LucasLLC

[pbontrager]

Stamping. Only note is to remember to update the Policy or the controller to remove old policies from the store once all generators are updated.

[Ritesh1905]

Curious: what is the reasoning for doing this at the learner and not at the generator? The trainer just pushes it weights and the generator can based on it's implementation (vLLM, sglang etc.) can modify the sd.

[pradeepfn]

This is an temp thing I did. It will be moved to generator sd loading + it will be moved out of the trainer/generator critical path based on efficiency numbers.

[pbontrager]

To add to Pradeep's answer, vLLM already handles it's own hf -> vllm mapping. The only reason we've recreated it is so we can add a shaded loading solution which we want to eventually upstream. It will be on the generator side like he said.