Skip to content

Fix device map in limited mem #1190

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jun 20, 2024
Merged

Fix device map in limited mem #1190

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
5b06cab
fix memory in limited env
tastelikefeet Jun 20, 2024
db660c5
fix
tastelikefeet Jun 20, 2024
58444eb
fix cpu
tastelikefeet Jun 20, 2024
98d46fb
Merge commit 'db660c5bfd6b87bad336425ddd2837ef2dc3c071' into fix/limi…
tastelikefeet Jun 20, 2024
6fda6c9
fix
tastelikefeet Jun 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .dev_scripts/dockerci.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@ MODELSCOPE_CACHE_DIR_IN_CONTAINER=/modelscope_cache
CODE_DIR=$PWD
CODE_DIR_IN_CONTAINER=/swift
echo "$USER"
gpus='0,1 2,3 4,5 6,7'
cpu_sets='0-15 16-31 32-47 48-63'
gpus='0,1 2,3'
cpu_sets='0-15 16-31'
cpu_sets_arr=($cpu_sets)
is_get_file_lock=false
CI_COMMAND=${CI_COMMAND:-bash .dev_scripts/ci_container_test.sh python tests/run.py --parallel 2 --run_config tests/run_config.yaml}
Expand Down
2 changes: 1 addition & 1 deletion swift/llm/infer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ def prepare_model_template(args: InferArguments,
else:
logger.info(f'device_count: {torch.cuda.device_count()}')
if device_map is None:
device_map = 'auto'
device_map = 'auto' if torch.cuda.device_count() > 1 else 'cuda:0'
if device_map == 'auto':
model_kwargs['low_cpu_mem_usage'] = True
model_kwargs['device_map'] = device_map
Expand Down
2 changes: 2 additions & 0 deletions swift/llm/rlhf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,8 @@ def llm_rlhf(args: RLHFArguments) -> Dict[str, Any]:
model_kwargs = {'low_cpu_mem_usage': True}
if is_dist() and not is_ddp_plus_mp():
model_kwargs['device_map'] = {'': local_rank}
elif torch.cuda.device_count() == 1:
model_kwargs['device_map'] = 'cuda:0'
else:
model_kwargs['device_map'] = 'auto'

Expand Down
2 changes: 2 additions & 0 deletions swift/llm/sft.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,8 @@ def llm_sft(args: SftArguments) -> Dict[str, Union[str, Any]]:
model_kwargs = {'low_cpu_mem_usage': True}
if is_dist() and not is_ddp_plus_mp():
model_kwargs['device_map'] = {'': local_rank}
elif torch.cuda.device_count() == 1:
model_kwargs['device_map'] = 'cuda:0'
elif not use_torchacc():
model_kwargs['device_map'] = 'auto'

Expand Down