Skip to content

Fix pad_input_tensors to pad the requested dim, not always dim 0 - #4150

Open
shoemoney wants to merge 1 commit into
huggingface:mainfrom
shoemoney:fix-pad-input-tensors-dim
Open

Fix pad_input_tensors to pad the requested dim, not always dim 0#4150
shoemoney wants to merge 1 commit into
huggingface:mainfrom
shoemoney:fix-pad-input-tensors-dim

Conversation

@shoemoney

Copy link
Copy Markdown

What does this PR do?

pad_input_tensors accepts a documented dim parameter, but the inner
_pad_input_tensors always resized new_size[0] while slicing on dim:

new_size[0] = batch_size + to_pad          # hardcodes dim 0
...
indices = tuple(slice(0, old_size[dim]) if i == dim else slice(None) for i in range(len(new_size)))
new_tensor[indices] = tensor               # but slices on dim

Any call with dim != 0 raises at the assignment:

input shape: (4, 6)
pad_input_tensors(t, batch_size=4, num_processes=9, dim=0)  # -> (9, 6), fine
pad_input_tensors(t, batch_size=4, num_processes=9, dim=1)  # -> RuntimeError:
    The expanded size of the tensor (9) must match the existing size (4)
    at non-singleton dimension 0.  Target sizes: [9, 6].  Tensor sizes: [4, 6]

The fix resizes new_size[dim] instead of new_size[0]. I also ported the
bounds-check and negative-dim normalization that the sibling
_pad_across_processes (~30 lines above) already has:

if dim >= len(tensor.shape) or dim < -len(tensor.shape):
    return tensor
if dim < 0:
    dim += len(tensor.shape)

I'm treating this as in-scope rather than scope creep: once dim is
actually used to index new_size, an out-of-range or negative dim needs
the same guard _pad_across_processes already applies, or it produces an
IndexError / silently-wrong new_size instead of the RuntimeError above.
It's the same "correctly handle dim" concern, not an unrelated addition,
and it keeps the two sibling pad functions behaving consistently for
invalid dim values.

Not fixed here

The to_pad divisibility arithmetic a few lines above this change (the
remainder / last_inputs / to_pad block) is a separate, pre-existing
bug and is already covered by #4093. This PR does not touch it — the diff
here is scoped to dim handling only.

Testing

Added test_slice_and_concatenate_dim to tests/test_utils.py, next to
the existing test_slice_and_concatenate. Confirmed it fails on main
and passes with the fix:

Before the fix:

tests/test_utils.py::UtilsTester::test_slice_and_concatenate_dim FAILED
...
E       RuntimeError: The expanded size of the tensor (9) must match the existing size (4) at non-singleton dimension 0.  Target sizes: [9, 6].  Tensor sizes: [4, 6]
src/accelerate/utils/operations.py:832: RuntimeError
=========================== short test summary info ============================
FAILED tests/test_utils.py::UtilsTester::test_slice_and_concatenate_dim - Run...
============================== 1 failed in 0.95s ===============================

After the fix:

tests/test_utils.py::UtilsTester::test_slice_and_concatenate_dim PASSED
tests/test_utils.py::UtilsTester::test_slice_and_concatenate PASSED
tests/test_utils.py::UtilsTester::test_pad_across_processes PASSED
=============================== 3 passed, 1 warning in 0.86s ===============================

Full tests/test_utils.py:

================== 46 passed, 4 skipped, 16 warnings in 8.64s ==================

ruff check and ruff format --check on the touched files both pass
clean (All checks passed! / 2 files already formatted).

Before submitting

  • Did you read the contributor guideline, Pull Request section?
  • Was this discussed/approved via a Github issue or the forum? (opening directly since it's a small, self-contained bug fix with a reproduction and a regression test)
  • Did you write any new necessary tests?

Who can review?

@SunMarc — core parts of the library.

_pad_input_tensors resized new_size[0] unconditionally while slicing on
the caller-supplied dim, so any dim != 0 raised a RuntimeError on the
tensor assignment. Resize new_size[dim] instead, and port the
bounds-check + negative-dim normalization already used by the sibling
_pad_across_processes so out-of-range/negative dim values are handled
the same way in both functions.

Adds a regression test covering dim=1 padding.
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