debug more7

vllm/patches/deepseek_v4.py

@@ -527,6 +527,49 @@ class DeepseekV4MegaMoEExperts(nn.Module):
         import os
         import nvfp4_megamoe_kernel as deep_gemm
 
+        # One-time state dump: model params vs checkpoint keys for layer 0 attn
+        if int(os.environ.get('MEGA_MOE_DEBUG', '0')) and not getattr(self, '_state_dump_done', False):
+            self._state_dump_done = True
+            from vllm.distributed import get_tensor_model_parallel_rank
+            if get_tensor_model_parallel_rank() == 0:
+                import gc
+                models = [o for o in gc.get_objects()
+                          if type(o).__name__ in ('DeepseekV4Model', 'DeepseekV4ForCausalLM')]
+                if models:
+                    m = models[0]
+                    sd = dict(m.named_parameters())
+                    l0 = sorted(k for k in sd if 'layers.0.' in k and ('attn' in k or 'self_attn' in k))
+                    print("=== MODEL params (layer 0 attn) ===")
+                    for k in l0:
+                        t = sd[k]
+                        nz = (t != 0).any().item()
+                        print(f"  {k}\n    shape={tuple(t.shape)} dtype={t.dtype} any_nonzero={nz}")
+
+                    from safetensors import safe_open
+                    import glob, json
+                    model_dir = '/model'
+                    idx_path = os.path.join(model_dir, 'model.safetensors.index.json')
+                    if os.path.exists(idx_path):
+                        with open(idx_path) as f:
+                            idx = json.load(f)
+                        l0_ckpt = sorted(k for k in idx['weight_map']
+                                         if 'layers.0.' in k and ('attn' in k or 'self_attn' in k))
+                        print(f"\n=== CHECKPOINT keys (layer 0 attn) ===")
+                        shard_of = idx['weight_map']
+                        shards_needed = sorted(set(shard_of[k] for k in l0_ckpt))
+                        shapes = {}
+                        for shard_name in shards_needed:
+                            shard_path = os.path.join(model_dir, shard_name)
+                            with safe_open(shard_path, framework='pt') as h:
+                                for k in l0_ckpt:
+                                    if shard_of[k] == shard_name:
+                                        shapes[k] = (tuple(h.get_tensor(k).shape), h.get_tensor(k).dtype)
+                        for k in l0_ckpt:
+                            s, d = shapes.get(k, ('?', '?'))
+                            print(f"  {k}\n    shape={s} dtype={d}")
+                    else:
+                        print(f"[state-dump] No index.json at {idx_path}")
+
         symm_buffer = self.get_symm_buffer()
         symm_buffer.experts_start_idx = self.experts_start_idx
         num_tokens = hidden_states.shape[0]
@@ -1329,7 +1372,30 @@ class DeepseekV4Model(nn.Module):
         # Pre-compute expert mapping ONCE.
         expert_mapping = self.get_expert_mapping()
 
+        # Debug: dump incoming checkpoint names for layer 0 attention o_* keys
+        if os.environ.get('MEGA_MOE_DEBUG'):
+            _o_sample = []
+            _o_seen = set()
+            for _n, _ in weights:
+                if 'layers.0.self_attn' in _n and 'o_' in _n:
+                    if _n not in _o_seen:
+                        _o_seen.add(_n)
+                        _o_sample.append(_n)
+                        if len(_o_sample) > 20:
+                            break
+            if _o_sample:
+                print(f"[LOAD-RAW] incoming layer 0 o_* checkpoint names:")
+                for _n in _o_sample:
+                    print(f"  {_n}")
+
         for name, loaded_weight in weights:
+            # Debug: trace o_b_proj/wo_b/o_a_proj/wo_a through the loader
+            if os.environ.get('MEGA_MOE_DEBUG') and ('o_b_proj' in name or 'wo_b' in name or 'o_a_proj' in name or 'wo_a' in name):
+                print(f"[LOAD-TRACE] candidate name={name!r} "
+                      f"in_params_dict={name in params_dict} "
+                      f"loaded_dtype={loaded_weight.dtype} "
+                      f"loaded_shape={tuple(loaded_weight.shape)}")
+
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 # Skip non-stacked layers and experts (experts handled below).
                 if ".experts." in name: