search

Lightning-AI / litgpt

Pretrain, finetune, deploy 20+ LLMs on your own data. Uses state-of-the-art techniques: flash attention, FSDP, 4-bit, LoRA, and more.
https://lightning.ai
Apache License 2.0
6.85k stars 726 forks source link

How to specify which GPU to use？ #1379

Closed Gooooooogo closed 2 weeks ago

Gooooooogo commented 2 weeks ago

def setup( checkpoint_dir: Path = Path("checkpoints/TinyLlama/TinyLlama-1.1B-Chat-v1.0"), out_dir: Path = Path("out/finetune/lora"), precision: Optional[str] = None, quantize: Optional[Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8-training"]] = None, devices: Union[int, str] = 1, lora_r: int = 8, lora_alpha: int = 16, lora_dropout: float = 0.05, lora_query: bool = True, lora_key: bool = False, lora_value: bool = True, lora_projection: bool = False, lora_mlp: bool = False, lora_head: bool = False, data: Optional[DataModule] = None,

I only found parameters specifying the number of GPUs, but how to specify which GPU to use?

rasbt commented 2 weeks ago

Good call, we should probably document this better. Say you have 4 GPUs, and you only want to use the first 2, you can do

CUDA_VISIBLE_DEVICES=0,1 litgpt finetune --checkpoint_dir checkpoints/microsoft/phi-2 --devices 2

or the 1st and last GPU:

CUDA_VISIBLE_DEVICES=0,3 litgpt finetune --checkpoint_dir checkpoints/microsoft/phi-2 --devices 2
Gooooooogo commented 2 weeks ago

Good call, we should probably document this better. Say you have 4 GPUs, and you only want to use the first 2, you can do

CUDA_VISIBLE_DEVICES=0,1 litgpt finetune --checkpoint_dir checkpoints/microsoft/phi-2 --devices 2

or the 1st and last GPU:

CUDA_VISIBLE_DEVICES=0,3 litgpt finetune --checkpoint_dir checkpoints/microsoft/phi-2 --devices 2

It works！ very thankful！