Skip to content

feat: force GPU usage with device_map to prevent CPU memory overflow#166

Open
Erlonidas wants to merge 1 commit intobytedance:masterfrom
Erlonidas:master
Open

feat: force GPU usage with device_map to prevent CPU memory overflow#166
Erlonidas wants to merge 1 commit intobytedance:masterfrom
Erlonidas:master

Conversation

@Erlonidas
Copy link
Copy Markdown

Problem

The original implementation was not properly utilizing GPU when available, causing the model to load on CPU RAM instead. This resulted in:

  • CPU RAM overflow: Maxing out 12GB system RAM
  • GPU underutilization: GPU sitting idle at 0GB usage despite being available
  • Significantly slower processing: Without GPU acceleration

Solution

This PR forces the model to load directly on GPU using device_map="cuda:0", ensuring proper GPU utilization from the start.

Changes

  • Load model directly on GPU with device_map="cuda:0" parameter
  • Add informative message showing GPU device name on startup
  • Use torch_dtype=torch.bfloat16 for optimized memory usage
  • Maintain CPU fallback for systems without GPU
  • Move device detection before model loading

Benefits

  • Prevents CPU RAM overflow: Model loads on GPU VRAM (~9GB) instead of system RAM
  • Better performance: Tested with 70 PDF pages processed in 2h40min with GPU vs much slower on CPU
  • Automatic GPU detection: No manual configuration needed
  • Memory efficient: Uses bfloat16 precision on GPU
  • User-friendly: Shows which GPU is being used at startup

Testing

Tested on Google Colab with Tesla T4 GPU (15GB VRAM):

  • ✅ Model successfully loads on GPU (~9GB VRAM usage)
  • ✅ CPU RAM stays low (~3GB instead of maxing out at 12GB)
  • ✅ 70 PDF pages processed successfully in 2h40min
  • ✅ Proper GPU utilization throughout processing

Before vs After

Before:

  • CPU RAM: 12.7GB / 12.7GB (maxed out) ❌
  • GPU VRAM: 0.0GB / 15.0GB (unused) ❌

After:

  • CPU RAM: ~3GB / 12.7GB (normal) ✅
  • GPU VRAM: ~9GB / 15.0GB (properly utilized) ✅

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant