From c96df0af5f7400bb632f1f2aacb7b599bc08af33 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Thu, 29 Oct 2026 23:45:56 -0700 Subject: [PATCH] [bugfix]: balance Flex calibration shards across resolutions --- scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py b/scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py index 233641bf7..a2b31e988 100644 --- a/scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py +++ b/scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py @@ -248,8 +248,12 @@ def _dp() -> tuple[int, int]: def _shard(rows: list[Any], shard: int, num_shards: int) -> list[Any]: + """This group's rows; striding a hash order keeps resolutions mixed (the plan alternates 480p / 768p groups).""" + import hashlib + dp_rank, dp_size = _dp() - return rows[shard * dp_size + dp_rank::num_shards * dp_size] + ordered = sorted(rows, key=lambda row: hashlib.sha256(row.id.encode()).hexdigest()) + return ordered[shard * dp_size + dp_rank::num_shards * dp_size] def _is_writer() -> bool: