Skip to content

Commit

Permalink
use iterator during weight loading
Browse files Browse the repository at this point in the history
  • Loading branch information
KeremTurgutlu committed Mar 12, 2024
1 parent 0b57d37 commit b3768f4
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion train.py
Original file line number Diff line number Diff line change
Expand Up @@ -622,7 +622,7 @@ def load_and_quantize_parallel(name_param, model, **kwargs):
n_workers = min(left, right)
if rank == 0 and args['verbose']:
print_func(f"Using n_workers: {n_workers} for loading")
parallel(load_and_quantize_parallel, weights.items(), n_workers=n_workers, threadpool=True,
parallel(load_and_quantize_parallel, iter(weights.items()), n_workers=n_workers, threadpool=True,
model=model, dtype=torch_dtype, device=local_rank, skip_names=load_param_skip_names,
is_meta_rank=(args["low_memory"] and rank!=0), verbose=args["verbose"], quant_method=quant_method)
if rank == 0 and args["verbose"]:
Expand Down

0 comments on commit b3768f4

Please sign in to comment.