[Python] Multi GPU Training [closed]

I'm training Llama-3.2-1B.

My Hardware Configuration:

GPUs:

NVIDIA GeForce RTX 3080

NVIDIA GeForce RTX 3080 Ti

NVIDIA GeForce RTX 3070

NVIDIA GeForce RTX 3070 Ti

RAM: 8 GB

Processor: Intel Core i5-8500


nvidia-smi output:

| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA GeForce RTX 3070 Off | 00000000:01:00.0 Off | N/A |
| 0% 31C P8 9W / 220W | 187MiB / 8192MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
| 1 NVIDIA GeForce RTX 3080 Off | 00000000:03:00.0 Off | N/A |
| 0% 36C P8 35W / 370W | 18MiB / 10240MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
| 2 NVIDIA GeForce RTX 3070 Ti Off | 00000000:04:00.0 Off | N/A |
| 0% 30C P8 8W / 290W | 18MiB / 8192MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
| 3 NVIDIA GeForce RTX 3080 Ti Off | 00000000:05:00.0 Off | N/A |
| 0% 32C P8 8W / 370W | 18MiB / 12288MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+


When I was training either my computer logs me our or gives me errors.

Errors Encountered:

[rank3]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 930.00 MiB. GPU 3 has a total capacity of 7.67 GiB of which 859.88 MiB is free. Process 2790 has 155.06 MiB memory in use. Including non-PyTorch memory, this process has 6.64 GiB memory in use. Of the allocated memory 5.73 GiB is allocated by PyTorch, and 608.34 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
(This specific error points to one of my 8GB cards, likely a 3070 or 3070 Ti).

[rank1]: RuntimeError: Input tensors need to be on the same GPU, but found the following tensor and device combinations:
[rank1]: [(torch.Size([2048, 2048]), device(type='cuda', index=0)), (torch.Size([128, 2048]), device(type='cuda', index=1)), (torch.Size([128, 2048]), device(type='cuda', index=0))]


What I'm Using:

Libraries: Deepspeed, LoRA, ZeRO Optimizer, Accelerate, quantization, fp16


This is My code:

MODEL_NAME = "meta-llama/Llama-3.2-1B"
TRAIN_PATH = "train.txt"
BATCH_SIZE = 1
NUM_WORKERS = 1
MAX_SEQUENCE_LENGTH = 128

# os.environ["CUDA_VISIBLE_DEVICES"] = "3,1,0,2"

class DatasetHelper(L.LightningDataModule):

def __init__(self,
model_name,
train_path,
batch_size,
num_workers,
max_sequence_length):
super().__init__()
self.model_name = model_name
self.train_path = train_path
self.batch_size = batch_size
self.num_workers = num_workers
self.max_sequence_length = max_sequence_length

self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_fast=False)
print("Tokenizer loaded")
self.tokenizer.pad_token = self.tokenizer.eos_token
self.tokenizer.padding_side = "right"
self.data_collator = DataCollatorForLanguageModeling(tokenizer=self.tokenizer, mlm=False)

def setup(self, stage: str):
if stage == "fit":
dataset_dict = load_dataset("text", data_files={"train": self.train_path})

def tokenize_function(examples):
return self.tokenizer(examples["text"], truncation=True, max_length=self.max_sequence_length)

self.tokenized_dataset = dataset_dict["train"].map(
tokenize_function,
batched=True,
num_proc=os.cpu_count(),
remove_columns=["text"],
)
print("Tokenized")

def train_dataloader(self):
print("Training dataloader...")
return DataLoader(
self.tokenized_dataset,
batch_size=self.batch_size,
collate_fn=self.data_collator,
num_workers=self.num_workers,
pin_memory=False,
)


class LlamaModelModule(L.LightningModule):
def __init__(self, model_name):
super().__init__()
print('Model Loading')
self.model = AutoModelForCausalLM.from_pretrained(model_name,
torch_dtype=float16,
load_in_8bit=True,
# quantization_config=self.bits_and_bytes_config(),
device_map=None)
self.model.gradient_checkpointing_enable()
print('Model Loaded')


def configure_optimizers(self):
return torch.optim.AdamW(self.model.parameters(), lr=5e-5)

def training_step(self, training_batch, batch_idx):
outputs = self.model(**training_batch)
loss = outputs.loss
self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True)
return loss

def bits_and_bytes_config(self):
return BitsAndBytesConfig(load_in_4bit=True)


def main():
dataset_helper = DatasetHelper(MODEL_NAME,
TRAIN_PATH,
BATCH_SIZE,
NUM_WORKERS,
MAX_SEQUENCE_LENGTH)

dataset_helper.setup("fit")
train_dataloader = dataset_helper.train_dataloader()

model = LlamaModelModule(MODEL_NAME)
torch.set_float32_matmul_precision('high')


trainer = L.Trainer(accelerator="gpu",
devices=-1,
strategy='deepspeed_stage_3',
precision='16-true',
accumulate_grad_batches=16,
enable_model_summary=True,
enable_progress_bar=True,
max_epochs=2)

trainer.fit(model, train_dataloader)

if __name__ == "__main__":
main()


and this is my ds_config.json:

{
"train_micro_batch_size_per_gpu": 2,
"zero_optimization": {
"stage": 3,
"offload_optimizer": {"device": "none"},
"offload_param": {"device": "none"},
"stage3_max_live_parameters": 3e8,
"stage3_prefetch_bucket_size": 2e8
},
"fp16": {"enabled": true},
"optimizer": {"type": "AdamW", "params": {"lr": 5e-5}}
}


My Questions:

1. Given my hardware (especially the 8GB VRAM limitation on some GPUs and the total 8GB system RAM), what are the most effective Deepspeed ZeRO stages and configurations I should be using to train Llama-3.2-1B?
2. Are there specific LoRA configurations (e.g., r, alpha, target_modules) that are more memory-efficient for this model size?
3. How can I better manage tensor placement and synchronization across my heterogeneous GPU setup to avoid the "All tensors must be on all devices" error? Are there specific accelerate configurations or Deepspeed settings that help with this?
4. Is my 8GB system RAM a significant bottleneck, even with ZeRO and Deepspeed offloading? Should I consider upgrading it?
5. Are there any other techniques or best practices I should implement to reduce VRAM usage or improve stability when training Llama-3.2-1B on this kind of setup?

Any guidance or specific configuration examples would be greatly appreciated!

Continue reading...