Enhance Model Loading By Providing Parallelism, Uses Optional Env Flag #36835
Conversation
|
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the |
There was a problem hiding this comment.
Very interesting!
Indeed we do this "sequentially" today. I think using an even bigger checkpoint like a 100B model or 400B would be nice.
Great splitting, having a separate function for load_shard_file!
- missing some API simplification:
multiprocessing.set_start_method("spawn", force=True)should be done by ourselves (because you should have access to the device you are loading on!)
Also needs to be guarded against TP (until stable / tested!)
src/transformers/modeling_utils.py
Outdated
| if json.loads(os.environ.get("HF_ENABLE_PARALLEL_LOADING", "false")): | ||
| num_workers = json.loads(os.environ.get("HF_PARALLEL_LOADING_WORKERS", "8")) | ||
| logger.info(f"Loading model weights in parallel with {num_workers} workers...") |
There was a problem hiding this comment.
this is nice BUT! I think we need some guards / good defaults:
- depends on the number of shard files
- should depend on the number of available threads
This will help us finding good sweetspots!
There was a problem hiding this comment.
Fully agree on point 1, we should just min() that on len(args_list).
Second point, I'd argue is an enhancement, for now the benefits are so great we can leave it to the user until itr 2.
(I'll be cooking this up real soon!)
src/transformers/modeling_utils.py
Outdated
| splits = tensor_name.split(".") | ||
| module = model_to_load | ||
|
|
||
| for split in splits[:-1]: | ||
| module = getattr(module, split) | ||
|
|
||
| last_key = splits.pop() |
There was a problem hiding this comment.
we have a util for that, see the tensor_parallel integration that also needs to get the module from the key name 😉
There was a problem hiding this comment.
Appreciate it! Will make the adjustment :)
ArthurZucker
added
Core: Modeling
|
Also, quick Q, is it normal for some of the CI tests to fail with the 4.50.0 branch, or is my code running amok? I checked the outputs and I don't think it's my code. Also, couldn't find a place where the model_utils.py run, I assumed anything with the Locally, wrt to just the tests for model_utils (including the parallel tests), they are passing, so I'm inclined to believe this code is stable, although it could certainly benefit from additional tests. It's setup to run optionally until we suss out any remaining edge cases. Also curious to hear what your thoughts are on how to improve the parallel tests, I can dream up some ideas, but I don't want to do anything that will result in waste. |
|
No longer needed, I've encapsulated this into the model loading parallelism if block Any ideas on where we could put this? I could run some experiments, but if you know of a great location where env vars are reduced into global configuration settings, that's probably where we want to do it, if there's some place where a bunch of package wide initialization occurs. |
…Y up the code by leveraging get_submodule.
…he multiprocessing start method are removed, now that the package handles this step internally.
|
@ArthurZucker Alright, I've made the requested changes, pls lmk if anything else needs touching up! Notable changes based on the feedback provided:
|
|
Hey! sorry I got caught up here and there will come back to this but no rush! |
|
We need to test on different settings, cluster no cluster, etc to make sure we are not introducing something that does not work well |
|
Np! More than appreciative you're assisting me with this :) Just give me a list of todo items and I'll execute, dying to merge this puppy! |
|
Hey Arthur, just bumping... Dying to make the changes needed to get this over the line |
There was a problem hiding this comment.
Hey @inf3rnus! Nice work, and super sorry for the wait on this one! You can see my review comments attached! 🤗
Basically, we need to be mostly careful about sharing the model to each child process. Since the recent changes, models are always loaded on meta (except with deepspeed when non-quantized, see comments), so it won't result in too much memory overhead, but the best would still be to be able to share the model between processes if possible. Most notably because the buffers are not on meta, so a model with some large buffers could see its memory requirements surge a lot with this (both on gpus devices and cpu, depending on the situation).
Also, could you make sure to rebase on latest main, and make sure that the functions you copy/pasted to be external functions are still similar to what is being done on main (because it won't be flagged as a merge conflict here) 🤗
src/transformers/modeling_utils.py
Outdated
|
|
||
| # We now update each layer of the meta model with the tensor module refs that were set to specific devices in the copy of the meta model for each worker | ||
| # We are transferring that state into the orginal ref (model_to_load) here | ||
| # This is required because model_to_load is pickled when using multiprocessing, which means the ref to model_to_load is different for each worker, so you only get some of the state with respect to the loaded tensors | ||
| # You could in theory return each worker's copy of the model and use .named_parameters(), and .named_buffers(), but this appears to be more robust | ||
| # in that all you have to care about are the names of the layers in the state dict, as long as the logic that lead to the creation of the state_dict is correct, this will also be correct | ||
| for state_dict_modules in state_dict_modules_list: | ||
| for full_name, param in state_dict_modules.items(): | ||
| *module_path, attr_name = full_name.split(".") | ||
| module_path = ".".join(module_path) | ||
| module = model_to_load.get_submodule(module_path) | ||
| setattr(module, attr_name, param) |
There was a problem hiding this comment.
Any way we could instead correctly share the model between processes here?
src/transformers/modeling_utils.py
Outdated
| # information we need in order to resolve all of the layers after multiprocessing which we write back to the original model_to_load meta model | ||
| state_dict_modules = resolve_state_dict_modules(model_to_load, state_dict, expected_keys) |
There was a problem hiding this comment.
We don't need to have this function here (it's very inefficient in case of NOT using multiprocessing, as it's basically useless) -> only the keys of the state_dict are enough information to then do it after the loading in the multiprocessing case
src/transformers/modeling_utils.py
Outdated
| # Use multiprocessing Pool for parallel execution, off by default | ||
| if json.loads(os.environ.get("HF_ENABLE_PARALLEL_LOADING", "false")): |
There was a problem hiding this comment.
This should be guarded against deepspeed here, as deepspeed is the only remaining path where the model in not on meta -> it will lead to exploding the memory as each process copy the model
|
Also, something very concerning to me is that in the current implementation, I think that if we have more state_dicts than num_workers, the second time a given worker starts processing a state_dict, the model it will copy already has some params loaded, which will blow up the memory very quickly (memory usage will be multiplied by a bit more than A simple workaround could be to deepcopy the model when instantiating the |
|
@Cyrilvallez Thanks Cyril, really appreciate your eyes on this 🙏 Moving through your feedback today, hope to have all concerns addressed by EOTD |
|
I've made your requested changes. I was able to greatly simplify things by just using threads to do the parallel loading. Which I when I created this PR was not possible. All tests pass for modeling_utils.py locally except for one ( Error for that is: You left this comment:
I believe this is solved by using threads now, but please lmk if there's something I'm missing. Please let me know if there is anything else you need from me, would love to get this merged asap! 🙏 |
|
Oh yeah, one other thing, I went to run |
src/transformers/modeling_utils.py
Outdated
| ): | ||
| map_location = torch.device([d for d in device_map.values() if d not in ["cpu", "disk"]][0]) | ||
| # Use multiprocessing Pool for parallel execution, off by default | ||
| if json.loads(os.environ.get("HF_ENABLE_PARALLEL_LOADING", "false")) and not is_deepspeed_zero3_enabled(): |
There was a problem hiding this comment.
do we have to use json loads here? a simple cast should work!
src/transformers/modeling_utils.py
Outdated
| with ThreadPoolExecutor(max_workers=num_workers) as executor: | ||
| with logging.tqdm(total=len(args_list), desc="Loading checkpoint shards") as pbar: | ||
| futures = [executor.submit(load_shard_file, arg) for arg in args_list] | ||
| for future in as_completed(futures): | ||
| result = future.result() | ||
| ( | ||
| _error_msgs, | ||
| disk_offload_index, | ||
| cpu_offload_index, | ||
| ) = result | ||
|
|
||
| # force memory release if loading multiple shards, to avoid having 2 state dicts in memory in next loop | ||
| del state_dict | ||
| error_msgs += _error_msgs | ||
|
|
||
| pbar.update(1) |
There was a problem hiding this comment.
can we put that in an external function?
Just to not make from_pretrained bigger than it already is!
| def is_true(value: Optional[str]) -> bool: | ||
| if value is None: | ||
| return False | ||
| return value.upper() in ENV_VARS_TRUE_VALUES | ||
|
|
There was a problem hiding this comment.
not a super fan of this, let's just do this:
os.environ.get("MY_VAR", "").upper() in ENV_VARS_TRUE_VALUES````
|
|
||
|
|
||
| # Declare the normal model_utils.py test as a sideffect of importing the module | ||
| from .test_modeling_utils import ModelUtilsTest # noqa |
There was a problem hiding this comment.
I don't think we need it to run on all models not even sure this tests it all?
…ABLE_PARALLEL_LOADING.
|
@Rocketknight1 bumping for visibility 🙏 |
|
Hey all, just checking in again, figure you guys are slammed. Very appreciative of the time you've spent reviewing this 🙏 |
|
@Cyrilvallez Tagging you again for a bump, it's right at the finish line. Very much appreciate everyone's eyes on this so far |
There was a problem hiding this comment.
ALright, sorry for the delay! This in indeed much cleaner, and I agree that using threads is much better as we share the objects by default. This should be thread safe as the different state dicts are orthogonal to each other, but did you check further by any chance? I.e. in the case of quantization, I did not check that we are not mutating other objects in ways that could go through race conditions, did you make sure?
Also, interested to see how much faster this is, given the GIL 🤔 Not entirely sure if the new thread would trigger new GPU streams, and if those would be run concurrently as well I must say, so hard to tell - did you benchmark a bit? 😁🤗
|
Oh also, please run |
|
Alright all changes requested have been made!
I just did a check on each of the quantizers, taking a look at You'd know best, but if we can guarantee that tying of weights only happens to the layers in a given shard, then we should be safe. Otherwise yeah, we're going to need to introduce locks. And it's really only one of the quantizers that I can see this happening in, it may be happening in the others, but if it is, I'm ignorant to it. Quantizer of concern: I don't think this should be a problem though because after looking at the code, if I'm not mistaken models always have loaded their input and output embeddings / encoder and decoder weights before loading the rest of the weights? It also only ties weights for the shard containing the module for the input embeddings. But outside of tying weights, I don't believe that there are any intersections in terms of assignment operations with respect to state, each quantization method only sets new weight values per module on a given layer via the state_dict that is constructed per thread, so each thread's set of layers of the model's overall state should be modified in isolation. They also all do this synchronously, so there should be no unexpected changes to a given module's state due to concurrent access.
def create_quantized_param(
self,
model: "PreTrainedModel",
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
state_dict: Dict[str, Any],
unexpected_keys: List[str],
):
"""
Each nn.Linear layer that needs to be quantized is processed here.
First, we set the value the weight tensor, then we move it to the target device. Finally, we quantize the module.
"""
if self.quantization_config.quant_type == "autoquant":
return
from torchao.quantization import quantize_
module, tensor_name = get_module_from_name(model, param_name)
if self.pre_quantized:
module._parameters[tensor_name] = torch.nn.Parameter(
param_value.to(device=target_device), requires_grad=param_value.requires_grad
)
if isinstance(module, nn.Linear):
module.extra_repr = types.MethodType(_linear_extra_repr, module)
else:
assert isinstance(self.quantization_config, TorchAoConfig)
module._parameters[tensor_name] = torch.nn.Parameter(
param_value, requires_grad=param_value.requires_grad
).to(device=target_device)
# if we are quantizing tied parameters, to avoid tying the quantized weights
# the correct order to do it is
# 1. load the weight to model
# 2. run tie_weights to populate the weights
# 3. quantize
input_embed = model.get_input_embeddings()
if self.quantization_config.untie_embedding_weights and id(module) == id(input_embed):
model.tie_weights()
setattr(model.config.get_text_config(decoder=True), "tie_word_embeddings", False)
# handle AOPerModuleConfig, introduced in torchao 0.11.0+
if self.quantization_config._get_ao_version() > version.Version("0.10.0"):
from torchao.quantization import AOPerModuleConfig
config = self.quantization_config.get_apply_tensor_subclass()
if isinstance(config, AOPerModuleConfig):
module_fqn, _ = param_name.rsplit(".", 1)
c = None
if module_fqn in config.module_fqn_to_config:
c = config.module_fqn_to_config[module_fqn]
else:
c = config.module_fqn_to_config.get("_default", None)
if c is not None:
# filter_fn: not filtering out any modules
quantize_(module, c, filter_fn=lambda x, fqn: True)
return
quantize_(module, self.quantization_config.get_apply_tensor_subclass())
Surprisingly, and maybe not so surprisingly the performance is even better. I did do a run on Presumably what's happening is a lot of this work's IO bound, the only major part AFAIK that isn't is the deserialization of the weights into system RAM. So while only one thread may deserialize weights, the rest are free to do work while that guy sits on the main thread. If the disk can provide the throughput and you have vCPUs at your disposal, you will see speed up benefits! The same can be said about offloading weights from RAM into the GPU, that's all likely handled by system calls and driver code that's divorced from any work that needs to actually happen within the python process itself. That's my theory at least, if we infinite time, we could certainly drill into it 😄... Prev results were around New results are around |
|
@Cyrilvallez @ArthurZucker Pinging again, I appreciate your patience 🙏 |
What does this PR do?
modeling_utils.pyhas been modified to allow for the parallelized loading of weights. Two env vars have been introduced,HF_ENABLE_PARALLEL_LOADINGwhich allows weights to be loaded in a multiprocessing pool andHF_PARALLEL_LOADING_WORKERSwhich specifies how many child processes you want to spawn to handle loading.Only works for sharded torch based weights and safe tensors for now, although the same pattern can be extended to other weights by other contributors.
This can significantly speed up loading large models on big hardware, ~50%.
e.g. facebook/opt-30b on an AWS EC2 g4dn.metal can be made to load in ~30s with this modification vs ~55s without it.
Using napkin math, this increase in efficiency should equate to thousands of dollars monthly if not more in cold boot compute expenses being saved globally for anyone using HF on the cloud.
Please let me know if anything else needs to be changed to move this forward!
Before submitting
Pull Request section?
documentation guidelines, and
here are tips on formatting docstrings.
@Rocketknight1 @gante