Skip to content

fix: enforce tensor packing limits and remove redundant clones #129

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yulangz wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: enforce tensor packing limits and remove redundant clones #129

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 47 additions & 1 deletion awex/tests/test_tensor_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from awex.util.tensor_util import (
group_tensors_by_shape_and_dtype,
reconstruct_tensors_from_groups,
release_tensors,
)

mp.set_start_method("spawn", force=True)
Expand Down Expand Up @@ -154,7 +155,52 @@ def test_max_tensor_size_limit(self):
# Check that each group size is within limit
for group in result_groups:
group_size = group.element_size() * group.numel()
assert group_size <= max_size * 2
assert group_size <= max_size

@pytest.mark.parametrize(
("sizes", "expected_bytes"),
[([4, 2], [16, 8]), ([2, 7, 1], [8, 28, 4]), ([3, 2, 1], [20, 4])],
)
def test_size_bound_and_oversized_tensor_isolation(self, sizes, expected_bytes):
tensors = [torch.arange(size, dtype=torch.float32) for size in sizes]
groups, metadata = group_tensors_by_shape_and_dtype(tensors, max_tensor_size=20)
assert [group.nbytes for group in groups] == expected_bytes
for original, rebuilt in zip(
tensors, reconstruct_tensors_from_groups(groups, metadata), strict=True
):
torch.testing.assert_close(rebuilt, original, rtol=0, atol=0)

@pytest.mark.parametrize("num_tensors", [1, 3])
@pytest.mark.parametrize("noncontiguous", [False, True])
def test_packed_storage_survives_source_release(self, num_tensors, noncontiguous):
tensors = [
torch.arange(12, dtype=torch.float32).reshape(3, 4)
for _ in range(num_tensors)
]
if noncontiguous:
tensors = [tensor.t() for tensor in tensors]
expected = [tensor.clone() for tensor in tensors]
groups, metadata = group_tensors_by_shape_and_dtype(tensors)
release_tensors(tensors)
for original, rebuilt in zip(
expected, reconstruct_tensors_from_groups(groups, metadata), strict=True
):
torch.testing.assert_close(rebuilt, original, rtol=0, atol=0)

@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_cuda_packing_peak_allocates_only_output_storage(self):
tensors = [torch.ones(1024 * 1024, device="cuda") for _ in range(2)]
torch.cuda.synchronize()
baseline = torch.cuda.memory_allocated()
torch.cuda.reset_peak_memory_stats()
groups, metadata = group_tensors_by_shape_and_dtype(tensors)
torch.cuda.synchronize()
output_bytes = sum(group.nbytes for group in groups)
assert torch.cuda.max_memory_allocated() - baseline == output_bytes
for original, rebuilt in zip(
tensors, reconstruct_tensors_from_groups(groups, metadata), strict=True
):
torch.testing.assert_close(rebuilt, original, rtol=0, atol=0)

def test_complex_tensor_shapes(self):
"""Test function with complex tensor shapes."""
Expand Down
67 changes: 23 additions & 44 deletions awex/util/tensor_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,7 @@ def group_tensors_by_shape_and_dtype(
tensors: list[torch.Tensor], max_tensor_size: int = 5 * 1024 * 1024 * 1024, **kwargs
) -> tuple[list[torch.Tensor], list[dict]]:
"""
Group tensors by shape and dtype, ensuring each group's total size is less than max_tensor_size.
Group tensors by dtype with at most max_tensor_size bytes per group.
If one tensor is too large, it will be put into a separate group.

Args:
Expand All @@ -196,7 +196,7 @@ def group_tensors_by_shape_and_dtype(

Returns:
Tuple containing:
- List of tensor groups, where each group is a list of tensors with same shape and dtype
- List of flattened tensor groups with the same dtype
- List of metadata dictionaries for reconstructing original tensors

Raises:
Expand Down Expand Up @@ -236,53 +236,16 @@ def group_tensors_by_shape_and_dtype(
current_group_indices = []
current_group_size = 0

for orig_idx, tensor in group_tensors:
tensor_size = tensor.element_size() * tensor.numel()
current_group.append(tensor)
current_group_indices.append(orig_idx)
current_group_size += tensor_size
# Check if this tensor can fit in current group
if current_group_size > max_tensor_size:
# Finalize current group and start new one
# Use clone() to ensure a copy so caller can safely release original tensors
concatenated = torch.cat(
[t.reshape(-1) for t in current_group], dim=0
).clone()
final_tensor_groups.append(concatenated)
# Record metadata for tensors in this group
offset_elements = 0
for i, group_tensor in enumerate(current_group):
group_tensor_elements = group_tensor.numel()
metadata.append(
{
"original_index": current_group_indices[i],
"shape": group_tensor.shape,
"dtype": group_tensor.dtype,
"group_index": len(final_tensor_groups) - 1,
"offset": offset_elements,
"size": group_tensor_elements,
}
)
offset_elements += group_tensor_elements
# Start new group
current_group = []
current_group_indices = []
current_group_size = 0

# Finalize any remaining group
if current_group:
# Use clone() to ensure a copy so caller can safely release original tensors
concatenated = torch.cat(
[t.reshape(-1) for t in current_group], dim=0
).clone()
def finalize_group(group: list[torch.Tensor], indices: list[int]) -> None:
# cat already owns independent storage, even with one input.
concatenated = torch.cat([t.reshape(-1) for t in group], dim=0)
final_tensor_groups.append(concatenated)
# Record metadata for tensors in this group
offset_elements = 0
for i, group_tensor in enumerate(current_group):
for original_index, group_tensor in zip(indices, group, strict=True):
group_tensor_elements = group_tensor.numel()
metadata.append(
{
"original_index": current_group_indices[i],
"original_index": original_index,
"shape": group_tensor.shape,
"dtype": group_tensor.dtype,
"group_index": len(final_tensor_groups) - 1,
Expand All @@ -291,6 +254,22 @@ def group_tensors_by_shape_and_dtype(
}
)
offset_elements += group_tensor_elements

for orig_idx, tensor in group_tensors:
tensor_size = tensor.element_size() * tensor.numel()
# Split before insertion. An individually oversized tensor forms
# its own group, rather than joining a preceding partial group.
if current_group and current_group_size + tensor_size > max_tensor_size:
finalize_group(current_group, current_group_indices)
current_group = []
current_group_indices = []
current_group_size = 0
current_group.append(tensor)
current_group_indices.append(orig_idx)
current_group_size += tensor_size

if current_group:
finalize_group(current_group, current_group_indices)
logger.info(
f"Grouped tensors, num groups: {len(final_tensor_groups)}, num tensors: {len(metadata)}"
)
Expand Down
2 changes: 2 additions & 0 deletions ci/test_cpu.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ set -euo pipefail
export AWEX_DEVICE_TYPE=cpu
python -m pytest --timeout=120 \
awex/tests/test_imports.py \
awex/tests/test_tensor_utils.py \
awex/tests/test_cuda_ipc.py \
awex/tests/test_runtime_contracts.py \
awex/tests/test_weights_reader_init.py \
awex/tests/test_meta_server.py \
Expand Down