Different between memory usage reported by driver and inside container

[thien-lm]

What happened: the memory consumption reported by driver and hami is different

What you expected to happen: reported memory consumption on both hami and driver are the same

How to reproduce it (as minimally and precisely as possible):
deploy gpt2 with vllm, each container consume 15k MiB of GPU DRAM
Anything else we need to know?:

• The output of nvidia-smi -a on vllm container:

• The output of nvidia-smi -a on host:

Environment:

• HAMi version: 2.6.0
• nvidia driver or other AI device driver version:
• Docker version from docker version
• Docker command, image and tag used
• Kernel version from uname -a
• Others:

[archlitchi]

yes, opensource hami only count the memory allocated by 'cuMalloc', so it will cause this problem

[thien-lm]

hi @archlitchi, with only one GPU, this gap is quite small and is acceptable, but with multi GPU and NCCL, this gap is quite big.
I thinks we should have mechanism to handle multi GPUs case.
I also tested and find out that hami also do not count memory occupied by cuCtxCreate.

[thien-lm]

@archlitchi , i've found out that NCCL call cuMemCreate many times at initialization step, so i modify cuMemCreate in libvgpu as this:

CUresult cuMemCreate ( CUmemGenericAllocationHandle* handle, size_t size, const CUmemAllocationProp* prop, unsigned long long flags ) {
    LOG_INFO("cuMemCreate:");
    ENSURE_RUNNING();
    CUdevice dev;
    cuCtxGetDevice(&dev);
    add_gpu_device_memory_usage(getpid(),dev,size,2);
    CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemCreate,handle,size,prop,flags);
    return res;
}

The memory gap is much smaller than before, from 2833 MiB to 239MiB, which acceptable for me. I want more insight about this approach but i am not good at C right now, hope you can check this solution :D

[thien-lm]

This is the code i tested:
package:
pip install torch transformers datasets accelerate trl bitsandbytes

code:

  import os
  os.environ["NCCL_CUMEM_ENABLE"] = "1"
  os.environ["LIBCUDA_LOG_LEVEL"] = "0"
  import torch
  from transformers import GPT2Tokenizer, GPT2LMHeadModel
  from torch.nn import CrossEntropyLoss
  
  model_name = "distilgpt2"
  tokenizer = GPT2Tokenizer.from_pretrained(model_name)
  tokenizer.pad_token = tokenizer.eos_token
  
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  n_gpus = torch.cuda.device_count()
  print(f"Using {n_gpus} GPUs")
  
  model = GPT2LMHeadModel.from_pretrained(model_name)
  if n_gpus > 1:
     model = torch.nn.DataParallel(model)
  model.to(device)
  
  text = "This is a sample sentence for SLM training. " * 50
  batch_size = 8
  
  inputs = tokenizer([text] * batch_size, return_tensors="pt", max_length=512, padding="max_length", truncation=True)
  input_ids = inputs["input_ids"].to(device)
  labels = input_ids.clone()
  
  optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
  loss_fct = CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)
  
  model.train()
  for epoch in range(100000):
      outputs = model(input_ids=input_ids)
      logits = outputs.logits
  
      shift_logits = logits[..., :-1, :].contiguous()
      shift_labels = labels[..., 1:].contiguous()
  
      loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
  
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
  
      print(f"Epoch {epoch+1} | Loss: {loss.item():.4f}")

[archlitchi]

that's great, have you submitted a PR related to that?

[thien-lm]

@archlitchi i've added it here Project-HAMi/HAMi-core#98