fix: resolve tensor file overwrite between target and draft models (#21694)
Co-authored-by: jiangguangya <jiangguangya@baidu.com> Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
This commit is contained in:
co-authored by
jiangguangya
Khoa Pham
parent
7ce2352323
commit
b8a2dcd300
@@ -1379,9 +1379,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
|||||||
f"mem usage={self.weight_load_mem_usage:.2f} GB."
|
f"mem usage={self.weight_load_mem_usage:.2f} GB."
|
||||||
)
|
)
|
||||||
if self.server_args.debug_tensor_dump_output_folder is not None:
|
if self.server_args.debug_tensor_dump_output_folder is not None:
|
||||||
|
dump_folder = self.server_args.debug_tensor_dump_output_folder
|
||||||
|
if self.spec_algorithm.is_eagle():
|
||||||
|
role = "draft" if self.is_draft_worker else "target"
|
||||||
|
dump_folder = os.path.join(dump_folder, role)
|
||||||
register_forward_hook_for_model(
|
register_forward_hook_for_model(
|
||||||
self.model,
|
self.model,
|
||||||
self.server_args.debug_tensor_dump_output_folder,
|
dump_folder,
|
||||||
self.server_args.debug_tensor_dump_layers,
|
self.server_args.debug_tensor_dump_layers,
|
||||||
self.tp_size,
|
self.tp_size,
|
||||||
self.tp_rank,
|
self.tp_rank,
|
||||||
|
|||||||
@@ -6254,7 +6254,11 @@ class ServerArgs:
|
|||||||
"--debug-tensor-dump-output-folder",
|
"--debug-tensor-dump-output-folder",
|
||||||
type=str,
|
type=str,
|
||||||
default=ServerArgs.debug_tensor_dump_output_folder,
|
default=ServerArgs.debug_tensor_dump_output_folder,
|
||||||
help="The output folder for dumping tensors.",
|
help=(
|
||||||
|
"The output folder for dumping tensors. "
|
||||||
|
"In Eagle mode, tensor outputs from draft and target models "
|
||||||
|
"are stored in separate subdirectories ('draft' and 'target')."
|
||||||
|
),
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--debug-tensor-dump-layers",
|
"--debug-tensor-dump-layers",
|
||||||
|
|||||||
Reference in New Issue
Block a user