AssertionError when using model "google/gemma-2b" with multi-gpus

[tritct]

System Info

Information

• [X] Docker
• [ ] The CLI directly

Tasks

• [X] An officially supported command
• [ ] My own modifications

Reproduction

I'm trying to run Docker on 2 A16 GPUS using model_id "google/gemma-2b". But after the model downloading step I run into AssertionError like the following.

Expected behavior

When I run with only 1 GPU it can initialize just fine. This issue only happen when I try to use multigpu.