Add `device_map_option` to `accelerate` args

bigscience-workshop / lm-evaluation-harness

A framework for few-shot evaluation of autoregressive language models.

MIT License

101 stars 30 forks source link

Closed jon-tow closed 2 years ago

jon-tow commented 2 years ago

Adds support for accelerate device_map options for finer grain control.
- This addresses the problem of running out of memory on the 0-th rank GPU because input tokens share the same memory space as some of the partitioned parameters - simply set device_map_options="balanced_low_0".