[ci] fix: dataloader in e2e ckpt test (#233)

Author: Luosuu

tests/checkpoints/test_trainer_saveload.py

@@ -133,7 +133,7 @@ def main():
     args.train.compute_train_steps(args.data.max_seq_len, args.data.train_size)
     train_dataloader = build_dataloader(
         dataset=train_dataset,
-        dataloader_type="streaming",
+        dataloader_type="native",
         micro_batch_size=args.train.micro_batch_size,
         global_batch_size=args.train.global_batch_size,
         dataloader_batch_size=args.train.dataloader_batch_size,
@@ -142,17 +142,14 @@ def main():
         rmpad=args.train.rmpad,
         rmpad_with_pos_ids=args.train.rmpad_with_pos_ids,
         bsz_warmup_ratio=args.train.bsz_warmup_ratio,
-        dyn_bsz_runtime=args.train.dyn_bsz_runtime,
+        bsz_warmup_init_mbtoken=args.train.bsz_warmup_init_mbtoken,
         dyn_bsz_margin=args.train.dyn_bsz_margin,
         dyn_bsz_buffer_size=args.train.dyn_bsz_buffer_size,
         collate_fn=None,
-        bsz_warmup_init_mbtoken=args.train.bsz_warmup_init_mbtoken,
-        infinity=True,
         num_workers=args.data.num_workers,
         drop_last=args.data.drop_last,
         pin_memory=args.data.pin_memory,
         prefetch_factor=args.data.prefetch_factor,
-        drop_resume_buffer=args.data.drop_resume_buffer,
     )
 
     logger.info_rank0("Prepare model")
@@ -351,7 +348,7 @@ def test_trainer_saveload_ep8():
         "--nnodes=1",
         "--nproc_per_node=8",
         "--master_port=4321",
-        "tests/utils/test_trainer_saveload.py",
+        "tests/checkpoints/test_trainer_saveload.py",
         "tests/checkpoints/ep8.yaml",
     ]
     ep8_result = subprocess.run(ep8_command, check=True)

veomni/checkpoint/dcp_checkpointer.py

@@ -175,8 +175,7 @@ def get_state_dict_with_ep_dim_preprocess(self, state_dict, action):
                 continue
 
             # each tensor in the state dict should only belong to one EP entry
-            if len(matches) > 1:
-                raise RuntimeError(f"Ambiguous EP spec match for state key '{name}': {matches}")
+            assert len(matches) == 1, f"Ambiguous EP spec match for state key '{name}': {matches}"
 
             ep_key = matches[0]
             cur_spec_info = ep_fqn2spec_info[ep_key]