[RFC]: In-place weights loading and model swapping

[22quinn]

Motivation.

We've implemented the initial version of inplace weights loading in #18745 and #19884.

Currently load_model will decide itself whether to load model from scratch or reload weighs inplace.
Should we give users more control? I've been thinking about making the inplace weights loading choice explicit.

In addition, a standalone weights reloading API makes it possible to support model swap in the original load_model method.

Proposed Change.

• collective_rpc("load_model")
  • this will always recreate model and load weights from sractch
• collective_rpc("update_config", overrides)
  • Update any config in model runner.
  • For example, collective_rpc("update_config", overrides={"load_config": {"load_format": "auto"}})
• collective_rpc("reload_weights")
  • this will explicitly reload weights inplace, without re-initialize the model
• Might need to move process_weights_after_loading from load_model to load_weights

Alternative approach
Instead of providing a standalone update_config API, we can embed the overrides in load_model and reload_weights APIs, in which we mutate the configs before performing model/weights loading.

Feedback Period.

No response

CC List.

cc @youkaichao @houseroad @yeqcharlotte @foreverlms

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[yeqcharlotte]

Wondering why do we need a dedicated update_load_config? Any downsides to add additional kwargs to reload_weights instead and handling load format change if there's any?

[22quinn]

If passing kwargs to reload_weights, there's a choice to make:

• Use it as a one-time loading config without mutating self.load_config, but feels weird as new weights are loaded but old config remains.
• Mutate self.load_config but it feels stealthy to do it in reload_weights?

Another consideration is evolving it to a dedicated update_config method, supporting all configs, helping cases like model swapping.