Merge pull request #679 from EleutherAI/fix-padding-ranks

[Refactor] Fix padding ranks

lm_eval/evaluator.py

@@ -191,6 +191,8 @@ def evaluate(
     samples = collections.defaultdict(list)
     requests = collections.defaultdict(list)
 
+    padding_requests = collections.defaultdict(int)
+
     # get lists of each type of request
     for task_name, task in task_dict.items():
         versions[task_name] = task.VERSION
@@ -239,6 +241,7 @@ def evaluate(
 
             # compute number of pseudobatches to pad with (FSDP/DDP require even batches among ranks)
             numpad = max(gathered_item) - gathered_item[lm.rank]
+            padding_requests[task.OUTPUT_TYPE] += numpad
 
     ### Run LM on inputs, get all outputs ###
     # execute each type of request
@@ -249,8 +252,8 @@ def evaluate(
         for req in reqs:
             cloned_reqs.extend([req] * req.repeats)
 
-        if (lm.world_size > 1) and (numpad > 0):
-            for _ in range(numpad):
+        if (lm.world_size > 1) and (padding_requests[reqtype] > 0):
+            for _ in range(padding_requests[reqtype]):
                 cloned_reqs.extend([req] * req.repeats)
 
         # run requests through model