chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,856 @@
|
||||
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.flex_checkpoint.dcp.key_validation import (
|
||||
AOAMappingEntry,
|
||||
AOASliceMapping,
|
||||
KeyValidationResult,
|
||||
ShapeMismatchInfo,
|
||||
_append_src_lines,
|
||||
_build_aoa_mappings,
|
||||
_classify_mappings,
|
||||
_describe_ops,
|
||||
_emit,
|
||||
_format_key_list,
|
||||
_format_pattern_groups,
|
||||
_format_slice_range,
|
||||
_get_signature,
|
||||
_group_by_signature,
|
||||
_group_keys_adaptive,
|
||||
_print_aoa_report,
|
||||
_print_standard_report,
|
||||
_slice_covers_full,
|
||||
_try_fold_src_keys,
|
||||
validate_and_report_keys_aoa,
|
||||
validate_and_report_keys_standard,
|
||||
)
|
||||
from paddle.distributed.flex_checkpoint.dcp.metadata import (
|
||||
LocalTensorIndex,
|
||||
LocalTensorMetadata,
|
||||
Metadata,
|
||||
)
|
||||
|
||||
|
||||
class TestSliceCoversFull(unittest.TestCase):
|
||||
def test_covers_full(self):
|
||||
sl = (slice(0, 4), slice(0, 8))
|
||||
self.assertTrue(_slice_covers_full(sl, (4, 8)))
|
||||
|
||||
def test_not_covers_partial(self):
|
||||
sl = (slice(0, 2), slice(0, 8))
|
||||
self.assertFalse(_slice_covers_full(sl, (4, 8)))
|
||||
|
||||
def test_not_covers_mismatched_dims(self):
|
||||
sl = (slice(0, 4),)
|
||||
self.assertFalse(_slice_covers_full(sl, (4, 8)))
|
||||
|
||||
def test_non_zero_start(self):
|
||||
sl = (slice(1, 4), slice(0, 8))
|
||||
self.assertFalse(_slice_covers_full(sl, (4, 8)))
|
||||
|
||||
|
||||
class TestFormatSliceRange(unittest.TestCase):
|
||||
def test_basic(self):
|
||||
src_sl = (slice(0, 4), slice(0, 8))
|
||||
dst_sl = (slice(0, 4), slice(0, 8))
|
||||
result = _format_slice_range(src_sl, dst_sl)
|
||||
self.assertIn("0:4", result)
|
||||
self.assertIn("0:8", result)
|
||||
self.assertIn("->", result)
|
||||
|
||||
def test_partial_slices(self):
|
||||
src_sl = (slice(2, 6),)
|
||||
dst_sl = (slice(0, 4),)
|
||||
result = _format_slice_range(src_sl, dst_sl)
|
||||
self.assertIn("2:6", result)
|
||||
self.assertIn("0:4", result)
|
||||
|
||||
|
||||
class TestTryFoldSrcKeys(unittest.TestCase):
|
||||
def test_fold_consecutive(self):
|
||||
keys = [f"model.experts.{i}.weight" for i in range(8)]
|
||||
result = _try_fold_src_keys(keys)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("{0..7}", result)
|
||||
|
||||
def test_no_fold_different_patterns(self):
|
||||
keys = ["model.a.weight", "model.b.weight"]
|
||||
result = _try_fold_src_keys(keys)
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_no_fold_multiple_varying_positions(self):
|
||||
keys = ["layer.0.expert.0.w", "layer.1.expert.1.w"]
|
||||
result = _try_fold_src_keys(keys)
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_single_key(self):
|
||||
result = _try_fold_src_keys(["a.0.b"])
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_empty(self):
|
||||
result = _try_fold_src_keys([])
|
||||
self.assertIsNone(result)
|
||||
|
||||
|
||||
class TestDescribeOps(unittest.TestCase):
|
||||
def test_single_with_permute(self):
|
||||
entry = AOAMappingEntry(
|
||||
dst_key="a.weight",
|
||||
dst_global_shape=(4, 8),
|
||||
slice_mappings=[
|
||||
AOASliceMapping(
|
||||
"b.weight",
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
["[1, 0]"],
|
||||
)
|
||||
],
|
||||
)
|
||||
result = _describe_ops(entry)
|
||||
self.assertIn("permute([1, 0])", result)
|
||||
|
||||
def test_concat_with_cast(self):
|
||||
entry = AOAMappingEntry(
|
||||
dst_key="a.weight",
|
||||
dst_global_shape=(8, 4),
|
||||
slice_mappings=[
|
||||
AOASliceMapping(
|
||||
"b.weight",
|
||||
(slice(0, 4), slice(0, 4)),
|
||||
(slice(0, 4), slice(0, 4)),
|
||||
["bfloat16"],
|
||||
),
|
||||
AOASliceMapping(
|
||||
"c.weight",
|
||||
(slice(0, 4), slice(0, 4)),
|
||||
(slice(4, 8), slice(0, 4)),
|
||||
["bfloat16"],
|
||||
),
|
||||
],
|
||||
)
|
||||
result = _describe_ops(entry)
|
||||
self.assertIn("concat", result)
|
||||
self.assertIn("cast(bfloat16)", result)
|
||||
|
||||
def test_no_ops(self):
|
||||
entry = AOAMappingEntry(
|
||||
dst_key="a.weight",
|
||||
dst_global_shape=(4, 8),
|
||||
slice_mappings=[
|
||||
AOASliceMapping(
|
||||
"b.weight",
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
None,
|
||||
)
|
||||
],
|
||||
)
|
||||
result = _describe_ops(entry)
|
||||
self.assertEqual(result, "")
|
||||
|
||||
def test_empty_slice_mappings(self):
|
||||
entry = AOAMappingEntry(
|
||||
dst_key="a.weight", dst_global_shape=(4,), slice_mappings=[]
|
||||
)
|
||||
result = _describe_ops(entry)
|
||||
self.assertEqual(result, "")
|
||||
|
||||
|
||||
class TestClassifyMappings(unittest.TestCase):
|
||||
def _make_entry(self, dst_key, src_key, pp=None, multi_src=False):
|
||||
if multi_src:
|
||||
sms = [
|
||||
AOASliceMapping(src_key, (slice(0, 4),), (slice(0, 4),), pp),
|
||||
AOASliceMapping(
|
||||
src_key + ".2", (slice(0, 4),), (slice(4, 8),), pp
|
||||
),
|
||||
]
|
||||
else:
|
||||
sms = [AOASliceMapping(src_key, (slice(0, 4),), (slice(0, 4),), pp)]
|
||||
return AOAMappingEntry(
|
||||
dst_key=dst_key, dst_global_shape=(8,), slice_mappings=sms
|
||||
)
|
||||
|
||||
def test_rename_only(self):
|
||||
entry = self._make_entry("model.layers.2.w", "model.layers.0.w")
|
||||
rename, transform, struct = _classify_mappings([entry])
|
||||
self.assertEqual(len(rename), 1)
|
||||
self.assertEqual(len(transform), 0)
|
||||
self.assertEqual(len(struct), 0)
|
||||
|
||||
def test_with_transform(self):
|
||||
entry = self._make_entry(
|
||||
"model.layers.2.w", "model.layers.0.w", ["[1, 0]"]
|
||||
)
|
||||
rename, transform, struct = _classify_mappings([entry])
|
||||
self.assertEqual(len(rename), 0)
|
||||
self.assertEqual(len(transform), 1)
|
||||
self.assertEqual(len(struct), 0)
|
||||
|
||||
def test_structural_multi_src(self):
|
||||
entry = self._make_entry(
|
||||
"model.layers.2.qkv", "model.layers.0.q", multi_src=True
|
||||
)
|
||||
rename, transform, struct = _classify_mappings([entry])
|
||||
self.assertEqual(len(rename), 0)
|
||||
self.assertEqual(len(transform), 0)
|
||||
self.assertEqual(len(struct), 1)
|
||||
|
||||
def test_structural_different_pattern(self):
|
||||
entry = self._make_entry("model.decoder.0.w", "model.encoder.0.w")
|
||||
rename, transform, struct = _classify_mappings([entry])
|
||||
self.assertEqual(len(struct), 1)
|
||||
|
||||
|
||||
class TestGroupBySignature(unittest.TestCase):
|
||||
def test_same_signature_grouped(self):
|
||||
entries = []
|
||||
for i in range(4):
|
||||
entries.append(
|
||||
AOAMappingEntry(
|
||||
dst_key=f"model.layers.{i}.w",
|
||||
dst_global_shape=(4,),
|
||||
slice_mappings=[
|
||||
AOASliceMapping(
|
||||
f"src.layers.{i}.w",
|
||||
(slice(0, 4),),
|
||||
(slice(0, 4),),
|
||||
["[1, 0]"],
|
||||
)
|
||||
],
|
||||
)
|
||||
)
|
||||
groups = _group_by_signature(entries)
|
||||
self.assertEqual(len(groups), 1)
|
||||
self.assertEqual(len(next(iter(groups.values()))), 4)
|
||||
|
||||
def test_different_signatures(self):
|
||||
e1 = AOAMappingEntry(
|
||||
dst_key="model.layers.0.w",
|
||||
dst_global_shape=(4,),
|
||||
slice_mappings=[
|
||||
AOASliceMapping(
|
||||
"src.layers.0.w", (slice(0, 4),), (slice(0, 4),), None
|
||||
)
|
||||
],
|
||||
)
|
||||
e2 = AOAMappingEntry(
|
||||
dst_key="model.layers.0.qkv",
|
||||
dst_global_shape=(12,),
|
||||
slice_mappings=[
|
||||
AOASliceMapping(
|
||||
"src.layers.0.q", (slice(0, 4),), (slice(0, 4),), None
|
||||
),
|
||||
AOASliceMapping(
|
||||
"src.layers.0.k", (slice(0, 4),), (slice(4, 8),), None
|
||||
),
|
||||
],
|
||||
)
|
||||
groups = _group_by_signature([e1, e2])
|
||||
self.assertEqual(len(groups), 2)
|
||||
|
||||
|
||||
class TestGetSignature(unittest.TestCase):
|
||||
def test_digits_normalized(self):
|
||||
entry = AOAMappingEntry(
|
||||
dst_key="model.layers.5.weight",
|
||||
dst_global_shape=(4,),
|
||||
slice_mappings=[
|
||||
AOASliceMapping(
|
||||
"src.layers.5.weight", (slice(0, 4),), (slice(0, 4),), None
|
||||
)
|
||||
],
|
||||
)
|
||||
sig = _get_signature(entry)
|
||||
self.assertIn("{N}", sig)
|
||||
self.assertNotIn("5", sig)
|
||||
|
||||
|
||||
class TestFormatPatternGroups(unittest.TestCase):
|
||||
def test_basic_output(self):
|
||||
entries = [
|
||||
AOAMappingEntry(
|
||||
dst_key="model.layers.0.w",
|
||||
dst_global_shape=(4, 8),
|
||||
slice_mappings=[
|
||||
AOASliceMapping(
|
||||
"src.layers.0.w",
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
["[1, 0]"],
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
groups = {"sig1": entries}
|
||||
lines, next_idx = _format_pattern_groups(groups, "test", 1)
|
||||
self.assertTrue(any("Pattern #1" in l for l in lines))
|
||||
self.assertEqual(next_idx, 2)
|
||||
|
||||
def test_numbering_continues(self):
|
||||
e1 = [
|
||||
AOAMappingEntry(
|
||||
dst_key="a.0.w",
|
||||
dst_global_shape=(4,),
|
||||
slice_mappings=[
|
||||
AOASliceMapping(
|
||||
"b.0.w", (slice(0, 4),), (slice(0, 4),), None
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
e2 = [
|
||||
AOAMappingEntry(
|
||||
dst_key="c.0.w",
|
||||
dst_global_shape=(4,),
|
||||
slice_mappings=[
|
||||
AOASliceMapping(
|
||||
"d.0.w", (slice(0, 4),), (slice(0, 4),), None
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
groups = {"sig1": e1, "sig2": e2}
|
||||
lines, next_idx = _format_pattern_groups(groups, "test", 5)
|
||||
self.assertEqual(next_idx, 7)
|
||||
|
||||
def test_max_patterns_truncation(self):
|
||||
import paddle.distributed.flex_checkpoint.dcp.key_validation as kv
|
||||
|
||||
old = kv._MAX_PATTERNS_SHOWN
|
||||
kv._MAX_PATTERNS_SHOWN = 2
|
||||
try:
|
||||
groups = {}
|
||||
for i in range(5):
|
||||
groups[f"sig{i}"] = [
|
||||
AOAMappingEntry(
|
||||
dst_key=f"x.{i}.w",
|
||||
dst_global_shape=(4,),
|
||||
slice_mappings=[
|
||||
AOASliceMapping(
|
||||
f"y.{i}.w", (slice(0, 4),), (slice(0, 4),), None
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
lines, _ = _format_pattern_groups(groups, "test", 1)
|
||||
self.assertTrue(any("more" in l for l in lines))
|
||||
finally:
|
||||
kv._MAX_PATTERNS_SHOWN = old
|
||||
|
||||
|
||||
class TestAppendSrcLines(unittest.TestCase):
|
||||
def test_few_srcs(self):
|
||||
sms = [
|
||||
AOASliceMapping("a.w", (slice(0, 4),), (slice(0, 4),), None),
|
||||
AOASliceMapping("b.w", (slice(0, 4),), (slice(4, 8),), None),
|
||||
]
|
||||
lines = []
|
||||
_append_src_lines(lines, sms)
|
||||
self.assertEqual(len(lines), 2)
|
||||
self.assertIn("SRC:", lines[0])
|
||||
self.assertIn("+", lines[1])
|
||||
|
||||
def test_many_srcs_foldable(self):
|
||||
sms = [
|
||||
AOASliceMapping(
|
||||
f"experts.{i}.w",
|
||||
(slice(0, 4),),
|
||||
(slice(i * 4, (i + 1) * 4),),
|
||||
None,
|
||||
)
|
||||
for i in range(10)
|
||||
]
|
||||
lines = []
|
||||
_append_src_lines(lines, sms)
|
||||
# Should fold into single line with ×N
|
||||
self.assertTrue(any("\u00d7" in l for l in lines))
|
||||
|
||||
def test_many_srcs_not_foldable(self):
|
||||
sms = [
|
||||
AOASliceMapping(
|
||||
"src_alpha.w", (slice(0, 4),), (slice(0, 4),), None
|
||||
),
|
||||
AOASliceMapping("src_beta.w", (slice(0, 4),), (slice(4, 8),), None),
|
||||
AOASliceMapping(
|
||||
"src_gamma.w", (slice(0, 4),), (slice(8, 12),), None
|
||||
),
|
||||
AOASliceMapping(
|
||||
"src_delta.w", (slice(0, 4),), (slice(12, 16),), None
|
||||
),
|
||||
AOASliceMapping(
|
||||
"src_epsilon.w", (slice(0, 4),), (slice(16, 20),), None
|
||||
),
|
||||
AOASliceMapping(
|
||||
"src_zeta.w", (slice(0, 4),), (slice(20, 24),), None
|
||||
),
|
||||
]
|
||||
lines = []
|
||||
_append_src_lines(lines, sms)
|
||||
# Should show first 2, ..., last 1
|
||||
self.assertTrue(any("more" in l for l in lines))
|
||||
|
||||
|
||||
class TestGroupKeysAdaptive(unittest.TestCase):
|
||||
def test_basic_grouping(self):
|
||||
keys = [
|
||||
"model.layers.0.weight",
|
||||
"model.layers.1.weight",
|
||||
"model.layers.2.weight",
|
||||
"model.embed.weight",
|
||||
]
|
||||
groups = _group_keys_adaptive(keys)
|
||||
self.assertEqual(len(groups), 2)
|
||||
# layers.* grouped together
|
||||
layer_group = [g for g in groups.values() if len(g) == 3]
|
||||
self.assertEqual(len(layer_group), 1)
|
||||
|
||||
def test_no_digits(self):
|
||||
keys = ["model.weight", "model.bias"]
|
||||
groups = _group_keys_adaptive(keys)
|
||||
self.assertEqual(len(groups), 2)
|
||||
|
||||
|
||||
class TestFormatKeyList(unittest.TestCase):
|
||||
def test_few_keys(self):
|
||||
keys = {"a.w", "b.w", "c.w"}
|
||||
lines = _format_key_list(keys)
|
||||
self.assertEqual(len(lines), 3)
|
||||
|
||||
def test_many_keys_grouped(self):
|
||||
keys = {f"model.layers.{i}.weight" for i in range(100)}
|
||||
lines = _format_key_list(keys)
|
||||
# Should be grouped and folded
|
||||
self.assertTrue(len(lines) < 100)
|
||||
self.assertTrue(any("[" in l for l in lines))
|
||||
|
||||
def test_empty(self):
|
||||
lines = _format_key_list(set())
|
||||
self.assertEqual(lines, [])
|
||||
|
||||
|
||||
class TestEmit(unittest.TestCase):
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation.logger")
|
||||
def test_normal_output(self, mock_logger):
|
||||
lines = ["line1", "line2", "line3"]
|
||||
_emit(lines)
|
||||
self.assertEqual(mock_logger.info.call_count, 3)
|
||||
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation.logger")
|
||||
def test_truncation(self, mock_logger):
|
||||
import paddle.distributed.flex_checkpoint.dcp.key_validation as kv
|
||||
|
||||
old_max = kv._MAX_TOTAL_LINES
|
||||
kv._MAX_TOTAL_LINES = 5
|
||||
try:
|
||||
lines = ["x"] * 20
|
||||
_emit(lines)
|
||||
# 5 lines + 1 truncation msg = 6
|
||||
self.assertEqual(mock_logger.info.call_count, 6)
|
||||
finally:
|
||||
kv._MAX_TOTAL_LINES = old_max
|
||||
|
||||
|
||||
class TestPrintStandardReport(unittest.TestCase):
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_all_matched(self, mock_emit):
|
||||
result = KeyValidationResult()
|
||||
_print_standard_report(result, "/tmp/ckpt", 100)
|
||||
lines = mock_emit.call_args[0][0]
|
||||
self.assertTrue(any("[OK]" in l for l in lines))
|
||||
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_with_missing_and_unexpected(self, mock_emit):
|
||||
result = KeyValidationResult(
|
||||
missing_keys={"a.w", "b.w"},
|
||||
unexpected_keys={"c.w"},
|
||||
shape_mismatches=[ShapeMismatchInfo("d.w", (4, 8), (4, 16))],
|
||||
)
|
||||
_print_standard_report(result, "/tmp/ckpt", 100)
|
||||
lines = mock_emit.call_args[0][0]
|
||||
self.assertTrue(any("Missing" in l for l in lines))
|
||||
self.assertTrue(any("Unexpected" in l for l in lines))
|
||||
self.assertTrue(any("Shape" in l for l in lines))
|
||||
self.assertTrue(any("Matched: 98/100" in l for l in lines))
|
||||
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_shape_mismatch_truncation(self, mock_emit):
|
||||
import paddle.distributed.flex_checkpoint.dcp.key_validation as kv
|
||||
|
||||
old = kv._MAX_SHAPE_MISMATCHES
|
||||
kv._MAX_SHAPE_MISMATCHES = 2
|
||||
try:
|
||||
mismatches = [
|
||||
ShapeMismatchInfo(f"k{i}", (4,), (8,)) for i in range(5)
|
||||
]
|
||||
result = KeyValidationResult(
|
||||
missing_keys={"x"}, shape_mismatches=mismatches
|
||||
)
|
||||
_print_standard_report(result, "/tmp/ckpt", 10)
|
||||
lines = mock_emit.call_args[0][0]
|
||||
self.assertTrue(any("and 3 more" in l for l in lines))
|
||||
finally:
|
||||
kv._MAX_SHAPE_MISMATCHES = old
|
||||
|
||||
|
||||
class TestPrintAoaReport(unittest.TestCase):
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_all_resolved(self, mock_emit):
|
||||
mappings = [
|
||||
AOAMappingEntry(
|
||||
dst_key="a.w",
|
||||
dst_global_shape=(4,),
|
||||
slice_mappings=[
|
||||
AOASliceMapping("b.w", (slice(0, 4),), (slice(0, 4),), None)
|
||||
],
|
||||
is_identity=False,
|
||||
),
|
||||
]
|
||||
result = KeyValidationResult()
|
||||
_print_aoa_report(result, mappings, set(), "/tmp/ckpt")
|
||||
lines = mock_emit.call_args[0][0]
|
||||
self.assertTrue(any("[OK]" in l for l in lines))
|
||||
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_with_missing(self, mock_emit):
|
||||
mappings = [
|
||||
AOAMappingEntry(
|
||||
"a.w",
|
||||
(4,),
|
||||
[AOASliceMapping("b.w", (slice(0, 4),), (slice(0, 4),), None)],
|
||||
)
|
||||
]
|
||||
result = KeyValidationResult(
|
||||
missing_keys={"c.w"}, unexpected_keys={"d.w"}
|
||||
)
|
||||
_print_aoa_report(result, mappings, {"removed.w"}, "/tmp/ckpt")
|
||||
lines = mock_emit.call_args[0][0]
|
||||
self.assertTrue(any("Missing" in l for l in lines))
|
||||
self.assertTrue(any("Unexpected" in l for l in lines))
|
||||
self.assertTrue(any("Removed" in l for l in lines))
|
||||
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_randomly_initialized_keys(self, mock_emit):
|
||||
mappings = []
|
||||
result = KeyValidationResult(
|
||||
randomly_initialized_keys={"init.w", "init.b"}
|
||||
)
|
||||
_print_aoa_report(result, mappings, set(), "/tmp/ckpt")
|
||||
lines = mock_emit.call_args[0][0]
|
||||
self.assertTrue(any("Initialized (2)" in l for l in lines))
|
||||
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_removed_keys_truncation(self, mock_emit):
|
||||
mappings = []
|
||||
removed = {f"removed.key.{i}" for i in range(10)}
|
||||
result = KeyValidationResult()
|
||||
_print_aoa_report(result, mappings, removed, "/tmp/ckpt")
|
||||
lines = mock_emit.call_args[0][0]
|
||||
self.assertTrue(any("more" in l for l in lines))
|
||||
|
||||
|
||||
class TestBuildAoaMappings(unittest.TestCase):
|
||||
def test_basic(self):
|
||||
engine = MagicMock()
|
||||
td1 = MagicMock()
|
||||
td1.shape = [4, 8]
|
||||
td1.slices = [
|
||||
(
|
||||
"src.w",
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
None,
|
||||
)
|
||||
]
|
||||
td2 = MagicMock()
|
||||
td2.shape = [8, 8]
|
||||
td2.slices = [
|
||||
(
|
||||
"src.q",
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
["[1, 0]"],
|
||||
),
|
||||
(
|
||||
"src.k",
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
(slice(4, 8), slice(0, 8)),
|
||||
["[1, 0]"],
|
||||
),
|
||||
]
|
||||
ov = MagicMock()
|
||||
ov.items.return_value = sorted({"dst.qkv": td2, "dst.w": td1}.items())
|
||||
engine.output_vars = ov
|
||||
|
||||
results = _build_aoa_mappings(engine)
|
||||
self.assertEqual(len(results), 2)
|
||||
qkv = next(r for r in results if r.dst_key == "dst.qkv")
|
||||
self.assertFalse(qkv.is_identity)
|
||||
self.assertEqual(len(qkv.slice_mappings), 2)
|
||||
|
||||
def test_identity_detection(self):
|
||||
engine = MagicMock()
|
||||
td = MagicMock()
|
||||
td.shape = [4, 8]
|
||||
td.slices = [
|
||||
(
|
||||
"same.key",
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
None,
|
||||
)
|
||||
]
|
||||
ov = MagicMock()
|
||||
ov.items.return_value = [("same.key", td)]
|
||||
engine.output_vars = ov
|
||||
|
||||
results = _build_aoa_mappings(engine)
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertTrue(results[0].is_identity)
|
||||
|
||||
def test_none_tensor_desc_skipped(self):
|
||||
engine = MagicMock()
|
||||
ov = MagicMock()
|
||||
ov.items.return_value = [("a", None), ("b", None)]
|
||||
engine.output_vars = ov
|
||||
results = _build_aoa_mappings(engine)
|
||||
self.assertEqual(len(results), 0)
|
||||
|
||||
|
||||
class TestValidateAndReportKeysStandard(unittest.TestCase):
|
||||
def _make_metadata(self, keys_shapes):
|
||||
"""keys_shapes: dict of {key: shape_tuple}"""
|
||||
storage_metadata = {}
|
||||
state_dict_metadata = {}
|
||||
for key, shape in keys_shapes.items():
|
||||
idx = LocalTensorIndex(
|
||||
tensor_key=key,
|
||||
global_offset=tuple([0] * len(shape)),
|
||||
replica_id=0,
|
||||
)
|
||||
storage_metadata[idx] = f"{key}.distcp"
|
||||
state_dict_metadata[key] = [
|
||||
LocalTensorMetadata(
|
||||
global_offset=tuple([0] * len(shape)),
|
||||
local_shape=shape,
|
||||
dtype="float32",
|
||||
global_shape=shape,
|
||||
)
|
||||
]
|
||||
return Metadata(
|
||||
state_dict_metadata=state_dict_metadata,
|
||||
storage_metadata=storage_metadata,
|
||||
)
|
||||
|
||||
@patch("paddle.distributed.get_rank", return_value=0)
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_all_match(self, mock_emit, mock_rank):
|
||||
metadata = self._make_metadata({"w1": (4, 8), "w2": (4, 8)})
|
||||
state_dict = {
|
||||
"w1": paddle.zeros([4, 8]),
|
||||
"w2": paddle.zeros([4, 8]),
|
||||
}
|
||||
result = validate_and_report_keys_standard(
|
||||
[metadata], {"w1", "w2"}, None, False, "/tmp/ckpt", state_dict
|
||||
)
|
||||
self.assertEqual(len(result.missing_keys), 0)
|
||||
self.assertEqual(len(result.unexpected_keys), 0)
|
||||
|
||||
@patch("paddle.distributed.get_rank", return_value=0)
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_missing_keys(self, mock_emit, mock_rank):
|
||||
metadata = self._make_metadata({"w1": (4,)})
|
||||
state_dict = {
|
||||
"w1": paddle.zeros([4]),
|
||||
"w2": paddle.zeros([4]),
|
||||
}
|
||||
result = validate_and_report_keys_standard(
|
||||
[metadata], {"w1", "w2"}, None, False, "/tmp/ckpt", state_dict
|
||||
)
|
||||
self.assertIn("w2", result.missing_keys)
|
||||
|
||||
@patch("paddle.distributed.get_rank", return_value=0)
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_unexpected_keys(self, mock_emit, mock_rank):
|
||||
metadata = self._make_metadata({"w1": (4,), "w2": (4,), "w3": (4,)})
|
||||
state_dict = {"w1": paddle.zeros([4])}
|
||||
result = validate_and_report_keys_standard(
|
||||
[metadata], {"w1"}, None, False, "/tmp/ckpt", state_dict
|
||||
)
|
||||
self.assertIn("w2", result.unexpected_keys)
|
||||
self.assertIn("w3", result.unexpected_keys)
|
||||
|
||||
@patch("paddle.distributed.get_rank", return_value=0)
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_shape_mismatch(self, mock_emit, mock_rank):
|
||||
metadata = self._make_metadata({"w1": (4, 8)})
|
||||
state_dict = {"w1": paddle.zeros([4, 16])}
|
||||
result = validate_and_report_keys_standard(
|
||||
[metadata], {"w1"}, None, False, "/tmp/ckpt", state_dict
|
||||
)
|
||||
self.assertEqual(len(result.shape_mismatches), 1)
|
||||
self.assertEqual(result.shape_mismatches[0].src_global_shape, (4, 8))
|
||||
self.assertEqual(result.shape_mismatches[0].dst_global_shape, (4, 16))
|
||||
|
||||
@patch("paddle.distributed.get_rank", return_value=0)
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_replica_id_filtered(self, mock_emit, mock_rank):
|
||||
"""Keys with replica_id != 0 should be filtered out."""
|
||||
storage_metadata = {
|
||||
LocalTensorIndex(
|
||||
tensor_key="w1", global_offset=(0,), replica_id=0
|
||||
): "f1",
|
||||
LocalTensorIndex(
|
||||
tensor_key="w2", global_offset=(0,), replica_id=1
|
||||
): "f2",
|
||||
}
|
||||
metadata = Metadata(
|
||||
state_dict_metadata={
|
||||
"w1": [LocalTensorMetadata((0,), (4,), "float32", (4,))]
|
||||
},
|
||||
storage_metadata=storage_metadata,
|
||||
)
|
||||
state_dict = {"w1": paddle.zeros([4])}
|
||||
result = validate_and_report_keys_standard(
|
||||
[metadata], {"w1"}, None, False, "/tmp/ckpt", state_dict
|
||||
)
|
||||
self.assertEqual(len(result.unexpected_keys), 0)
|
||||
|
||||
@patch(
|
||||
"paddle.distributed.flex_checkpoint.dcp.key_validation._get_rank",
|
||||
return_value=1,
|
||||
)
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
@patch("paddle.distributed.all_gather_object")
|
||||
def test_non_rank0_no_print(self, mock_gather, mock_emit, mock_rank):
|
||||
metadata = self._make_metadata({"w1": (4,)})
|
||||
state_dict = {"w1": paddle.zeros([4])}
|
||||
|
||||
def gather_side_effect(out_list, obj, group=None):
|
||||
out_list.clear()
|
||||
out_list.append(obj)
|
||||
|
||||
mock_gather.side_effect = gather_side_effect
|
||||
validate_and_report_keys_standard(
|
||||
[metadata], {"w1"}, None, True, "/tmp/ckpt", state_dict
|
||||
)
|
||||
mock_emit.assert_not_called()
|
||||
|
||||
|
||||
class TestValidateAndReportKeysAoa(unittest.TestCase):
|
||||
def _make_mock_engine(self):
|
||||
engine = MagicMock()
|
||||
td1 = MagicMock()
|
||||
td1.shape = [4, 8]
|
||||
td1.slices = [
|
||||
(
|
||||
"src.w1",
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
None,
|
||||
)
|
||||
]
|
||||
td2 = MagicMock()
|
||||
td2.shape = [8, 8]
|
||||
td2.slices = [
|
||||
(
|
||||
"src.q",
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
["[1, 0]"],
|
||||
),
|
||||
(
|
||||
"src.k",
|
||||
(slice(0, 4), slice(0, 8)),
|
||||
(slice(4, 8), slice(0, 8)),
|
||||
["[1, 0]"],
|
||||
),
|
||||
]
|
||||
# output_vars: need .items() for _build_aoa_mappings and iteration for values()
|
||||
ov = MagicMock()
|
||||
ov.items.return_value = sorted({"dst.w1": td1, "dst.qkv": td2}.items())
|
||||
ov.values.return_value = [td1, td2]
|
||||
ov.__iter__ = lambda self: iter({"dst.w1": td1, "dst.qkv": td2})
|
||||
ov.__getitem__ = lambda self, k: {"dst.w1": td1, "dst.qkv": td2}[k]
|
||||
engine.output_vars = ov
|
||||
engine.need_add_output_vars = ["dst.init"]
|
||||
engine.need_remove_input_vars = ["src.removed"]
|
||||
engine.input_vars = MagicMock()
|
||||
engine.input_vars.keys.return_value = [
|
||||
"src.w1",
|
||||
"src.q",
|
||||
"src.k",
|
||||
"src.removed",
|
||||
"src.leftover",
|
||||
]
|
||||
engine.context = MagicMock()
|
||||
engine.context.get_all_dst_state_keys.return_value = {
|
||||
"dst.w1",
|
||||
"dst.qkv",
|
||||
"dst.init",
|
||||
}
|
||||
return engine
|
||||
|
||||
@patch("paddle.distributed.get_rank", return_value=0)
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_all_resolved(self, mock_emit, mock_rank):
|
||||
engine = self._make_mock_engine()
|
||||
metadata = MagicMock()
|
||||
result = validate_and_report_keys_aoa(engine, metadata, "/tmp/ckpt")
|
||||
# dst.w1 and dst.qkv are covered; dst.init is randomly initialized
|
||||
self.assertEqual(len(result.missing_keys), 0)
|
||||
# src.leftover not consumed and not removed
|
||||
self.assertIn("src.leftover", result.unexpected_keys)
|
||||
self.assertIn("dst.init", result.randomly_initialized_keys)
|
||||
|
||||
@patch("paddle.distributed.get_rank", return_value=0)
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_truly_missing(self, mock_emit, mock_rank):
|
||||
engine = self._make_mock_engine()
|
||||
# Add a dst key that is NOT covered
|
||||
engine.context.get_all_dst_state_keys = lambda: {
|
||||
"dst.w1",
|
||||
"dst.qkv",
|
||||
"dst.init",
|
||||
"dst.missing",
|
||||
}
|
||||
metadata = MagicMock()
|
||||
result = validate_and_report_keys_aoa(engine, metadata, "/tmp/ckpt")
|
||||
self.assertIn("dst.missing", result.missing_keys)
|
||||
|
||||
@patch("paddle.distributed.get_rank", return_value=1)
|
||||
@patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
|
||||
def test_non_rank0_no_print(self, mock_emit, mock_rank):
|
||||
engine = self._make_mock_engine()
|
||||
metadata = MagicMock()
|
||||
validate_and_report_keys_aoa(engine, metadata, "/tmp/ckpt")
|
||||
mock_emit.assert_not_called()
|
||||
|
||||
|
||||
class TestColorHelpers(unittest.TestCase):
|
||||
def test_no_color(self):
|
||||
from paddle.distributed.flex_checkpoint.dcp.key_validation import _C
|
||||
|
||||
self.assertEqual(_C.green("test"), "test")
|
||||
self.assertEqual(_C.yellow("test"), "test")
|
||||
self.assertEqual(_C.red("test"), "test")
|
||||
self.assertEqual(_C.cyan("test"), "test")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user