fix(speechlm2): harden conversion and index reuse

@@ -75,6 +75,10 @@ def _adapt_strategy_for_conversion_world(strategy_cfg: dict, world_size: int) ->         return strategy_cfg      conversion_dp_size = world_size // non_dp_size+    # Training recipes commonly pin the full data-parallel world explicitly.+    # Conversion may intentionally use fewer ranks, so carrying that value over+    # would construct a mesh larger than the conversion process group.+    strategy_cfg["dp_size"] = conversion_dp_size     replicate_size = int(strategy_cfg.get("dp_replicate_size") or 1)     if replicate_size > 1 and (conversion_dp_size % replicate_size != 0 or replicate_size >= conversion_dp_size):         strategy_cfg["dp_replicate_size"] = 1

@@ -1220,6 +1220,23 @@ def maybe_log_moe_metrics(self):         if not layer_loads:             return +        # Routing counts are tiny, but the metric helpers perform reductions+        # and dtype conversions on the tensors' current device. At the end of+        # a packed training step the CUDA allocator can have effectively no+        # headroom, so even ``load.mean()`` may fail. The distributed+        # all-reduce above must happen on CUDA; after it completes, move the+        # detached metric payload to CPU before doing any reporting math.+        # Copy to CPU before converting dtype so ``.float()`` in the helper+        # cannot allocate a temporary CUDA tensor.+        layer_loads = {+            name: {+                **data,+                "expert_load": data["expert_load"].detach().to(device="cpu"),+                "aux_loss": (data["aux_loss"].detach().to(device="cpu") if data.get("aux_loss") is not None else None),+            }+            for name, data in layer_loads.items()+        }+         mode = moe_metrics_cfg.get("mode", "brief")         top_k = moe_metrics_cfg.get("top_k_experts", 5) @@ -1232,6 +1249,11 @@ def maybe_log_moe_metrics(self):         else:             metrics = compute_brief_metrics(layer_loads, top_k=top_k) +        # Lightning converts Python numbers to tensors on ``self.device``.+        # Keep the reporting path CPU-only after the all-reduce by supplying+        # explicit CPU scalar tensors to ``log_dict``.+        metrics = {name: torch.as_tensor(value, device="cpu", dtype=torch.float32) for name, value in metrics.items()}+         # ``batch_size=1`` is required when training_step uses the         # ``dataloader_iter`` flavor: Lightning cannot infer the batch size         # from the closure, and these MoE metrics are model-internal

@@ -111,6 +111,10 @@ _MANIFEST_REUSE_BATCH_BYTES = 64 << 20  +class NativeTarRouteSignatureMismatch(ValueError):+    """The authenticated source route does not describe the target manifest rows."""++ @dataclass(frozen=True) class NativeTarOrdinalMapSpec:     manifest_source_spec: object@@ -766,7 +770,7 @@ def _compare_native_tar_route_signatures(             zip_longest(source_rows, target_rows, fillvalue=missing)         ):             if source_data is missing or target_data is missing:-                raise ValueError(+                raise NativeTarRouteSignatureMismatch(                     f"Native-tar route reuse row count changed at shard {shard_index}: "                     f"source_route_rows={route_rows}, mismatch_at={row_index}"                 )@@ -777,7 +781,7 @@ def _compare_native_tar_route_signatures(                 target_data, path=target_manifest_path, row_index=row_index             )             if source_signature != target_signature:-                raise ValueError(+                raise NativeTarRouteSignatureMismatch(                     f"Native-tar routing signature changed at shard={shard_index} row={row_index}: "                     f"source={source_signature!r}, target={target_signature!r}, "                     f"source_manifest={source_manifest_path!r}, target_manifest={target_manifest_path!r}"@@ -887,7 +891,7 @@ def _initialize_native_tar_route_reuse_worker(source_pack_path: str) -> None: def _reuse_native_tar_ordinal_map_with_pack(     source_pack: IndexPack,     task: _NativeTarRouteReuseTask,-) -> tuple[int, IndexPackRecordValidationSummary]:+) -> tuple[int, IndexPackRecordValidationSummary | None, str | None]:     source_manifest_key = IndexPackCollectionSpec(         role="manifest",         kind=JSONL,@@ -896,20 +900,28 @@ def _reuse_native_tar_ordinal_map_with_pack(     ).key     source_manifest = source_pack.collection(source_manifest_key)     source_route = source_pack.collection(task.source_map.key)-    summary = _compare_native_tar_route_signatures(-        source_manifest,-        source_route,-        task.source_map,-        task.target_map,-        indexes_root=task.indexes_root,-    )+    try:+        summary = _compare_native_tar_route_signatures(+            source_manifest,+            source_route,+            task.source_map,+            task.target_map,+            indexes_root=task.indexes_root,+        )+    except NativeTarRouteSignatureMismatch as error:+        # Signature comparison completes before any route payload is copied, so+        # this target can safely fall back to a fresh build without retaining a+        # partial reused array. All other exceptions remain fatal.+        if any(path.exists() for path in task.output_paths):+            raise RuntimeError("Native-tar route mismatch left a partial reused array") from error+        return task.map_index, None, str(error)     _copy_packed_array_shards(source_route, task.output_paths)-    return task.map_index, summary+    return task.map_index, summary, None   def _reuse_native_tar_ordinal_map_worker(     task: _NativeTarRouteReuseTask,-) -> tuple[int, IndexPackRecordValidationSummary]:+) -> tuple[int, IndexPackRecordValidationSummary | None, str | None]:     if _NATIVE_TAR_ROUTE_REUSE_WORKER_PACK is None:         raise RuntimeError("Native-tar route reuse worker was not initialized")     return _reuse_native_tar_ordinal_map_with_pack(_NATIVE_TAR_ROUTE_REUSE_WORKER_PACK, task)@@ -979,12 +991,12 @@ def _reuse_native_tar_ordinal_array_specs(      temporary_directory = Path(temporary_directory)     reused: dict[bytes, IndexPackArraySpec] = {}-    unmatched = []     summaries = []     validated_manifest_keys: set[bytes] = set()     validators: list[Callable[[], None]] = []     tasks: list[_NativeTarRouteReuseTask] = []     task_metadata = {}+    successful_map_indices: set[int] = set()     with IndexPack(source_pack_path) as source_pack:         _authenticate_native_tar_route_source_pack(             source_pack,@@ -995,12 +1007,10 @@ def _reuse_native_tar_ordinal_array_specs(         for map_index, target_map in enumerate(target_maps):             source_map = source_by_tar_paths.get(target_map.tar_paths)             if source_map is None:-                unmatched.append(target_map)                 continue             if source_map.aggregate_manifest or target_map.aggregate_manifest:                 # Aggregate maps contain a second row-to-tar-shard route. Rebuild                 # both routes together instead of partially reusing only the member map.-                unmatched.append(target_map)                 continue              target_manifest_index_paths = tuple(@@ -1031,7 +1041,7 @@ def _reuse_native_tar_ordinal_array_specs(                 temporary_directory / f"reused-native-tar-route-{map_index:06d}-{shard_index:06d}.u32"                 for shard_index in range(target_map.sequence_count)             )-            reused[target_map.key] = IndexPackArraySpec(+            reused_spec = IndexPackArraySpec(                 role=target_map.role,                 kind=target_map.kind,                 source_spec=target_map.source_spec,@@ -1052,16 +1062,30 @@ def _reuse_native_tar_ordinal_array_specs(                 source_map,                 target_manifest_key,                 snapshot,+                reused_spec,             )          results = _iter_native_tar_route_reuse_results(source_pack, tasks, workers=native_tar_route_workers)-        for map_index, summary in results:-            target_map, source_map, target_manifest_key, snapshot = task_metadata[map_index]+        for map_index, summary, mismatch_reason in results:+            target_map, source_map, target_manifest_key, snapshot, reused_spec = task_metadata[map_index]+            if mismatch_reason is not None:+                assert summary is None+                logging.info(+                    "Rebuilding native-tar ordinal map after authenticated route mismatch: "+                    "target_key=%s source_key=%s reason=%s",+                    target_map.key.hex(),+                    source_map.key.hex(),+                    mismatch_reason,+                )+                continue+            assert summary is not None             if target_manifest_key not in validated_manifest_keys:                 validated_manifest_keys.add(target_manifest_key)                 summaries.append(summary)             snapshot.validate()             validators.append(snapshot.validate)+            reused[target_map.key] = reused_spec+            successful_map_indices.add(map_index)             logging.info(                 "Reused authenticated native-tar ordinal map: target_key=%s source_key=%s shards=%d rows=%d",                 target_map.key.hex(),@@ -1070,9 +1094,13 @@ def _reuse_native_tar_ordinal_array_specs(                 summary.records_checked,             ) +    maps_to_build = [+        target_map for map_index, target_map in enumerate(target_maps) if map_index not in successful_map_indices+    ]+     return (         reused,-        unmatched,+        maps_to_build,         (_sum_validation_summaries(*summaries) if summaries else _empty_validation_summary()),         validated_manifest_keys,         validators,

@@ -825,25 +825,33 @@ def test_converter_parallel_reuses_authenticated_routes(tmp_path):         + "\n"     )     create_jsonl_index(changed_manifest)-    failed_output = tmp_path / "parallel-failure.idxpack"-    failed_result = CliRunner().invoke(+    fallback_output = tmp_path / "parallel-fallback.idxpack"+    fallback_result = CliRunner().invoke(         main,         [             "--output",-            str(failed_output),+            str(fallback_output),             "--native-tar-route-workers",             "2",             *_native_tar_route_reuse_args(source_cfg, source_pack),             str(target_cfg),         ],     )-    assert failed_result.exit_code != 0-    assert "Native-tar routing signature changed" in failed_result.output-    assert not failed_output.exists()-    assert not list(tmp_path.glob(".parallel-failure.idxpack.native-tar-route.*"))+    assert fallback_result.exit_code == 0, fallback_result.output+    assert "native_tar_routes_reused=1" in fallback_result.output+    assert "native_tar_routes_built=1" in fallback_result.output+    with IndexPack(fallback_output) as pack:+        first_route = pack.collection(+            nemo_tar_ordinal_map_collection_key(str(expected_routes[0][0]), str(expected_routes[0][1]))+        )+        rebuilt_route = pack.collection(+            nemo_tar_ordinal_map_collection_key(str(changed_manifest), str(expected_routes[1][1]))+        )+        assert [first_route.value(index) for index in range(2)] == [1, 0]+        assert [rebuilt_route.value(index) for index in range(2)] == [0, 1]  -def test_converter_route_reuse_rejects_changed_ordered_routing_signature(tmp_path):+def test_converter_route_reuse_rebuilds_changed_ordered_routing_signature_without_copy(tmp_path, monkeypatch):     source_root = tmp_path / "source"     source_root.mkdir()     _, tar_path, source_cfg = _make_native_tar_routing_dataset(@@ -869,6 +877,11 @@ def test_converter_route_reuse_rejects_changed_ordered_routing_signature(tmp_pat         )     )     output = tmp_path / "target.idxpack"+    monkeypatch.setattr(+        converter,+        "_copy_packed_array_shards",+        lambda *_args, **_kwargs: pytest.fail("mismatched route payload must not be copied"),+    )      result = CliRunner().invoke(         main,@@ -880,11 +893,104 @@ def test_converter_route_reuse_rejects_changed_ordered_routing_signature(tmp_pat         ],     ) +    assert result.exit_code == 0, result.output+    assert "native_tar_routes_reused=0" in result.output+    assert "native_tar_routes_built=1" in result.output+    with IndexPack(output) as pack:+        route = pack.collection(nemo_tar_ordinal_map_collection_key(str(target_manifest), str(tar_path)))+        assert [route.value(index) for index in range(2)] == [0, 1]+++@pytest.mark.parametrize(+    ("target_audio_paths", "expected_route"),+    [+        (("A.wav",), [0]),+        (("A.wav", "B.wav", "C.wav"), [0, 1, 2]),+    ],+)+def test_converter_route_reuse_rebuilds_changed_row_count_without_copy(+    tmp_path, monkeypatch, target_audio_paths, expected_route+):+    source_root = tmp_path / "source"+    source_root.mkdir()+    _, tar_path, source_cfg = _make_native_tar_routing_dataset(+        source_root,+        [{"audio_filepath": name} for name in ("B.wav", "A.wav")],+        [("A.wav", b"A"), ("B.wav", b"B"), ("C.wav", b"C")],+    )+    source_pack = tmp_path / "source.idxpack"+    source_result = CliRunner().invoke(main, ["--output", str(source_pack), str(source_cfg)])+    assert source_result.exit_code == 0, source_result.output++    target_manifest = tmp_path / "target-manifest.jsonl"+    target_manifest.write_text("".join(json.dumps({"audio_filepath": name}) + "\n" for name in target_audio_paths))+    create_jsonl_index(target_manifest)+    target_cfg = tmp_path / "target.yaml"+    target_cfg.write_text(+        yaml.safe_dump(+            {+                "type": "nemo_tarred",+                "manifest_filepath": str(target_manifest),+                "tarred_audio_filepaths": str(tar_path),+            }+        )+    )+    monkeypatch.setattr(+        converter,+        "_copy_packed_array_shards",+        lambda *_args, **_kwargs: pytest.fail("mismatched route payload must not be copied"),+    )+    output = tmp_path / "target.idxpack"++    result = CliRunner().invoke(+        main,+        [+            "--output",+            str(output),+            *_native_tar_route_reuse_args(source_cfg, source_pack),+            str(target_cfg),+        ],+    )++    assert result.exit_code == 0, result.output+    assert "native_tar_routes_reused=0" in result.output+    assert "native_tar_routes_built=1" in result.output+    with IndexPack(output) as pack:+        route = pack.collection(nemo_tar_ordinal_map_collection_key(str(target_manifest), str(tar_path)))+        assert [route.value(index) for index in range(len(expected_route))] == expected_route+++def test_converter_route_reuse_unrelated_copy_error_remains_fatal(tmp_path, monkeypatch):+    _, _, source_cfg = _make_native_tar_routing_dataset(+        tmp_path,+        [{"audio_filepath": "B.wav"}, {"audio_filepath": "A.wav"}],+        [("A.wav", b"A"), ("B.wav", b"B")],+    )+    source_pack = tmp_path / "source.idxpack"+    source_result = CliRunner().invoke(main, ["--output", str(source_pack), str(source_cfg)])+    assert source_result.exit_code == 0, source_result.output++    def fail_after_partial_copy(_collection, output_paths):+        output_paths[0].write_bytes(b"partial")+        raise RuntimeError("unrelated route-copy failure")++    monkeypatch.setattr(converter, "_copy_packed_array_shards", fail_after_partial_copy)+    output = tmp_path / "target.idxpack"+    result = CliRunner().invoke(+        main,+        [+            "--output",+            str(output),+            *_native_tar_route_reuse_args(source_cfg, source_pack),+            str(source_cfg),+        ],+    )+     assert result.exit_code != 0-    assert "Native-tar routing signature changed at shard=0 row=0" in result.output+    assert isinstance(result.exception, RuntimeError)+    assert "unrelated route-copy failure" in str(result.exception)     assert not output.exists()     assert not list(tmp_path.glob(".target.idxpack.native-tar-route.*"))-    assert not list(tmp_path.glob(".target.idxpack.record-validation.*"))   def test_converter_route_reuse_builds_only_unmatched_tar_paths(tmp_path, monkeypatch):

@@ -63,7 +63,17 @@ def test_moe_metrics_skip_collective_between_intervals(monkeypatch): def test_moe_metrics_use_global_step_for_collection_and_detail_cadence(monkeypatch, global_step, expected_metric):     import nemo_automodel.components.moe.load_balance_metrics as metrics -    collect = MagicMock(return_value={"layer": torch.tensor([1.0])})+    expert_load = torch.tensor([1, 2], dtype=torch.int64, requires_grad=False)+    aux_loss = torch.tensor(0.25, requires_grad=True)+    collect = MagicMock(+        return_value={+            "layer": {+                "expert_load": expert_load,+                "aux_loss": aux_loss,+                "n_experts": 2,+            }+        }+    )     brief = MagicMock(return_value={"moe/mode": 0.0})     detailed = MagicMock(return_value={"moe/mode": 1.0})     monkeypatch.setattr(metrics, "collect_expert_loads", collect)@@ -77,10 +87,23 @@ def test_moe_metrics_use_global_step_for_collection_and_detail_cadence(monkeypat     if expected_metric == "brief":         brief.assert_called_once()         detailed.assert_not_called()+        payload = brief.call_args.args[0]     else:         detailed.assert_called_once()         brief.assert_not_called()+        payload = detailed.call_args.args[0]+    assert payload["layer"]["expert_load"].device.type == "cpu"+    assert payload["layer"]["expert_load"].dtype == torch.int64+    assert not payload["layer"]["expert_load"].requires_grad+    assert payload["layer"]["aux_loss"].device.type == "cpu"+    assert not payload["layer"]["aux_loss"].requires_grad+    assert payload["layer"]["n_experts"] == 2     model.log_dict.assert_called_once()+    logged = model.log_dict.call_args.args[0]+    assert logged+    assert all(isinstance(value, torch.Tensor) for value in logged.values())+    assert all(value.device.type == "cpu" for value in logged.values())+    assert all(value.dtype == torch.float32 for value in logged.values())   def test_moe_metrics_reject_nonpositive_interval():

@@ -901,6 +901,22 @@ def test_adapt_strategy_collapses_incompatible_hsdp_replicate_axis() -> None:     assert original["dp_replicate_size"] == 16  +def test_adapt_strategy_remaps_explicit_training_dp_size() -> None:+    original = {+        "dp_size": 256,+        "dp_replicate_size": 2,+        "tp_size": 1,+        "pp_size": 1,+        "cp_size": 1,+        "ep_size": 8,+    }+    adapted = to_hf._adapt_strategy_for_conversion_world(original, world_size=8)+    assert adapted["dp_size"] == 8+    assert adapted["dp_replicate_size"] == 2+    assert adapted["ep_size"] == 8+    assert original["dp_size"] == 256++ def test_adapt_strategy_preserves_compatible_hsdp_replicate_axis() -> None:     original = {         "dp_size": None,