@@ -590,7 +590,7 @@ def _sharded_state_dict(self, *args, **kwargs) -> "ShardedStateDict":
     # An MTP head left out of quantization is exempt from that: there is no quantization
     # error to recover there, and its CE materialises an fp32 [seq, vocab] tensor.
     skip_mtp_loss = _mtp_excluded_from_quantization(model)
-    if skip_mtp_loss:
+    if distill_cfg.skip_lm_loss and skip_mtp_loss:
         # Freeze the untrained MTP head: DDP's overlapped grad reduce asserts on params with no grad.
         for name, param in model.named_parameters():
             if "mtp" in name.split("."):
 
@@ -333,11 +333,6 @@ def save_pretrained(
         # medusa_heads and eagle_module only exist in the last stage.
         is_last_stage_main_rank = pp_rank == pp_size - 1 and tp_rank == 0
         is_writer_rank = self._is_sidecar_writer_rank(is_last_stage_main_rank)
-        # One writer per pipeline stage: other TP / DP / EP ranks hold the same layers (EP>1 ranks
-        # hold no gathered experts at all), and writing them too would race on the same files.
-        writes_layers = (
-            tp_rank == 0 and get_data_parallel_rank() == 0 and get_expert_model_parallel_rank() == 0
-        )
 
         quantization_format = self._get_quantization_format(self.model)
         if self._any_rank_uses_iq_quantization():
@@ -357,6 +352,22 @@ def save_pretrained(
                     "Megatron IQ1_S/IQ2_XS unified export currently requires pipeline model "
                     "parallel size 1"
                 )
+        # SequentialMLP rules index experts by local position, so EP>1 would collide. Agreed across
+        # ranks so that stages without such experts don't block in the collectives below.
+        if (
+            torch.distributed.is_initialized()
+            and get_expert_model_parallel_world_size() > 1
+            and self._any_rank(any(hasattr(m, "local_experts") for m in self.model.modules()))
+        ):
+            raise NotImplementedError(
+                "Export at expert parallel size > 1 needs grouped-GEMM experts; "
+                "export SequentialMLP (--no_moe_grouped_gemm) checkpoints at EP=1."
+            )
+        # One writer per pipeline stage: other TP / DP / EP ranks hold the same layers (EP>1 ranks
+        # hold no gathered experts at all), and writing them too would race on the same files.
+        writes_layers = (
+            tp_rank == 0 and get_data_parallel_rank() == 0 and get_expert_model_parallel_rank() == 0
+        )
 
         # Main export process
         layer_state_dicts = self.layer_state_dicts
@@ -800,15 +811,6 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
                             layer.mlp.shared_experts.gate_weight, layer_id, is_mtp=is_mtp
                         )
                 if hasattr(layer.mlp.experts, "local_experts"):
-                    # SequentialMLP rules index experts by local position, so EP>1 would collide.
-                    if (
-                        torch.distributed.is_initialized()
-                        and get_expert_model_parallel_world_size() > 1
-                    ):
-                        raise NotImplementedError(
-                            "Export at expert parallel size > 1 needs grouped-GEMM experts; "
-                            "export SequentialMLP (--no_moe_grouped_gemm) checkpoints at EP=1."
-                        )
                     if not self.rules.get("use_packed_local_experts", False):
                         for expert_id, expert in enumerate(layer.mlp.experts.local_experts):
                             self.rules["local_experts.linear_fc1"](
@@ -1319,11 +1321,15 @@ def _any_rank_uses_iq_quantization(self) -> bool:
         holding no IQ layer would skip the raise and then block in the next collective while its
         peers exit. Agree across ranks first, mirroring ``_gather_exclude_modules``.
         """
-        local_uses_iq = uses_iq_quantization(self.model)
+        return self._any_rank(uses_iq_quantization(self.model))
+
+    @staticmethod
+    def _any_rank(local: bool) -> bool:
+        """Whether ``local`` holds on any rank, so callers can raise everywhere or nowhere."""
         if not torch.distributed.is_initialized():
-            return local_uses_iq
+            return local
         per_rank = [None] * torch.distributed.get_world_size()
-        torch.distributed.all_gather_object(per_rank, local_uses_iq)
+        torch.distributed.all_gather_object(per_rank, local)
         return any(per_rank)
 
     def _get_quantization_format(self, module: torch.nn.Module):
 
@@ -216,7 +216,7 @@ def _test_topk_logits_kl_loss(top_k, rank, size):
     loss["kd_loss"].backward()
 
 
-def _test_skip_lm_loss_with_mtp(quantized, rank, size):
+def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, rank, size):
     """Test that skip_lm_loss only zeroes the main LM head, unless MTP is left out of quantization."""
     set_seed(SEED)
 
@@ -264,7 +264,7 @@ def _test_skip_lm_loss_with_mtp(quantized, rank, size):
         student_model.decoder.weight_quantizer = TensorQuantizer()
 
     distill_cfg = setup_distillation_config(
-        config_or_path=DistillationConfig(skip_lm_loss=True),
+        config_or_path=DistillationConfig(skip_lm_loss=skip_lm_loss),
         student_cfg=student_model.config,
         teacher_cfg=teacher_model.config,
     )
@@ -277,10 +277,11 @@ def _test_skip_lm_loss_with_mtp(quantized, rank, size):
     adjust_distillation_model_for_mcore(distillation_model, distill_cfg)
 
     # An untrained MTP head is frozen so DDP doesn't wait for its grads.
+    mtp_skipped = quantized and skip_lm_loss
     with distillation_model.hide_teacher_model():
         mtp_params = [p for n, p in distillation_model.named_parameters() if "mtp" in n.split(".")]
     assert mtp_params
-    assert all(p.requires_grad != quantized for p in mtp_params)
+    assert all(p.requires_grad != mtp_skipped for p in mtp_params)
 
     # Intercept each call to compute_language_model_loss and record return values.
     recorded_losses = []
@@ -314,8 +315,8 @@ def _recording_loss(labels, logits):
         f"Expected {mtp_num_layers + 1} loss calls, got {len(recorded_losses)}"
     )
     for i, loss in enumerate(recorded_losses[:-1]):
-        assert loss.any() != quantized, f"MTP head {i} loss should be zero iff MTP is skipped"
-    assert not recorded_losses[-1].any(), "Main LM head loss should be zero with skip_lm_loss=True"
+        assert loss.any() != mtp_skipped, f"MTP head {i} loss should be zero iff MTP is skipped"
+    assert recorded_losses[-1].any() != skip_lm_loss, "Main LM head loss is zero iff skip_lm_loss"
 
 
 def test_logits_kl_loss(dist_workers):
@@ -328,10 +329,12 @@ def test_topk_logits_kl_loss(dist_workers, top_k: int = 5):
     dist_workers.run(partial(_test_topk_logits_kl_loss, top_k))
 
 
-@pytest.mark.parametrize("quantized", [False, True])
-def test_skip_lm_loss_with_mtp(dist_workers, quantized):
+@pytest.mark.parametrize(
+    ("quantized", "skip_lm_loss"), [(False, True), (True, True), (True, False)]
+)
+def test_skip_lm_loss_with_mtp(dist_workers, quantized, skip_lm_loss):
     """Test that skip_lm_loss only zeroes the main LM head, unless MTP is left out of quantization."""
-    dist_workers.run(partial(_test_skip_lm_loss_with_mtp, quantized))
+    dist_workers.run(partial(_test_skip_lm_loss_with_mtp, quantized, skip_lm_loss))
 
 
 def test_mtp_excluded_from_quantization():
 
@@ -30,6 +30,7 @@
 from _test_utils.torch.megatron.models import get_mcore_gpt_model, get_mcore_hybrid_model
 from _test_utils.torch.megatron.utils import get_forward
 from _test_utils.torch.transformers_models import (
+    create_tiny_deepseek_v3_dir,
     create_tiny_llama_dir,
     create_tiny_nemotron_dir,
     create_tiny_nemotron_h_dir,
@@ -394,6 +395,51 @@ def test_megatron_iq_export_rejects_tensor_parallelism(qformat):
         exporter.save_pretrained("unused", "unused")
 
 
+def test_megatron_export_rejects_sequential_experts_at_ep():
+    """SequentialMLP experts are numbered by local position, so EP>1 export must refuse them."""
+    experts = torch.nn.Module()
+    experts.local_experts = torch.nn.ModuleList()
+    exporter = object.__new__(GPTModelExporter)
+    exporter.model = torch.nn.Sequential(experts)
+
+    with (
+        patch.object(exporter, "_is_sidecar_writer_rank", return_value=False),
+        patch.object(GPTModelExporter, "_any_rank", staticmethod(lambda local: local)),
+        patch.object(uem.torch.distributed, "is_initialized", return_value=True),
+        patch.object(uem, "get_expert_model_parallel_world_size", return_value=2),
+        patch.object(uem, "get_pipeline_model_parallel_rank", return_value=0),
+        patch.object(uem, "get_pipeline_model_parallel_world_size", return_value=1),
+        patch.object(uem, "get_tensor_model_parallel_rank", return_value=0),
+        pytest.raises(NotImplementedError, match="grouped-GEMM experts"),
+    ):
+        exporter.save_pretrained("unused", "unused")
+
+
+def _test_mla_export_keeps_hf_head_dim(model_dir, rank, size):
+    model = get_mcore_gpt_model(
+        tensor_model_parallel_size=size,
+        pipeline_model_parallel_size=1,
+        initialize_megatron=True,
+        num_layers=2,
+        hidden_size=128,
+        num_attention_heads=2,
+        vocab_size=128,
+        max_sequence_length=128,
+        transformer_impl="transformer_engine",
+        multi_latent_attention=True,
+    ).cuda()
+    hf_head_dim = transformers.AutoConfig.from_pretrained(model_dir).head_dim
+    # kv_channels is MLA's V head dim, which the exporter used to write as HF's head_dim.
+    assert model.config.kv_channels != hf_head_dim
+    exporter = GPTModelExporter(model, str(model_dir))
+    assert exporter._hf_text_config.head_dim == hf_head_dim
+
+
+def test_mla_export_keeps_hf_head_dim(dist_workers_size_1, tmp_path):
+    model_dir = create_tiny_deepseek_v3_dir(tmp_path)
+    dist_workers_size_1.run(partial(_test_mla_export_keeps_hf_head_dim, model_dir))
+
+
 @pytest.mark.parametrize("qformat", IQ_FORMAT_NAMES)
 def test_megatron_iq_export_rejects_pipeline_parallelism(qformat):
     """IQ packing requires PP=1 so the fused-MoE rejection reaches every rank.
@@ -424,6 +470,27 @@ def test_megatron_iq_export_rejects_pipeline_parallelism(qformat):
         exporter.save_pretrained("unused", "unused")
 
 
+def _verify_exported_metadata(export_dir: Path, model_type, quant_config, extra_module):
+    # The router bias decides expert routing: FP32 like Megatron's buffer and HF's checkpoints.
+    bias_dtypes = set()
+    for shard in export_dir.glob("*.safetensors"):
+        with safe_open(str(shard), framework="pt") as f:
+            for key in f.keys():  # noqa: SIM118 - safe_open is not iterable
+                if key.endswith(("expert_bias", "e_score_correction_bias")):
+                    bias_dtypes.add(f.get_slice(key).get_dtype())
+    assert bias_dtypes <= {"F32"}, bias_dtypes
+    if model_type == "qwen3_moe":
+        assert bias_dtypes
+    if model_type == "qwen3vl" and quant_config:
+        # The vision tower is copied through unquantized.
+        quant = json.loads((export_dir / "hf_quant_config.json").read_text())["quantization"]
+        assert "model.visual*" in quant["exclude_modules"]
+    if model_type == "llama" and quant_config is None and extra_module is None:
+        # The source generation config (top_p without do_sample) is copied through as-is.
+        generation_config = json.loads((export_dir / "generation_config.json").read_text())
+        assert generation_config["top_p"] == 0.5
+
+
 def _test_unified_export_megatron(
     tmp_path,
     model_type,
@@ -585,6 +652,9 @@ def _test_unified_export_megatron(
     if quant_config:
         _verify_model_quant_config(tmp_export_dir, quant_config, kv_cache_quant_cfg)
 
+    if rank == 0:
+        _verify_exported_metadata(tmp_export_dir, model_type, quant_config, extra_module)
+
     if rank == 0 and extra_module is None:
         # Names / shapes only: these Megatron weights are random, not loaded from model_dir.
         allow_missing = ()
@@ -650,6 +720,10 @@ def test_unified_export_megatron(
 ):
     if model_type == "llama":
         model_dir = create_tiny_llama_dir(tmp_path)
+        # Valid to load, but rejected by GenerationConfig.save_pretrained's validation.
+        (model_dir / "generation_config.json").write_text(
+            json.dumps({"top_p": 0.5, "do_sample": False})
+        )
     elif model_type == "qwen3vl":
         model_dir = create_tiny_qwen3vl_dir(tmp_path)
     elif model_type == "nemotron":