[webui] add seed controls for reproducibility (#10629)

Co-authored-by: Karunanidhi Mishra <11963379+kmishra1204@users.noreply.github.com>

Author: kmishra1204

src/llamafactory/chat/hf_engine.py

@@ -19,7 +19,7 @@
 from typing import TYPE_CHECKING, Any, Optional, Union
 
 import torch
-from transformers import GenerationConfig, TextIteratorStreamer
+from transformers import GenerationConfig, TextIteratorStreamer, set_seed
 from typing_extensions import override
 
 from ..data import get_template_and_fix_tokenizer
@@ -128,6 +128,7 @@ def _process_args(
         skip_special_tokens: Optional[bool] = input_kwargs.pop("skip_special_tokens", None)
         max_length: Optional[int] = input_kwargs.pop("max_length", None)
         max_new_tokens: Optional[int] = input_kwargs.pop("max_new_tokens", None)
+        seed: Optional[int] = input_kwargs.pop("seed", None)
         stop: Optional[Union[str, list[str]]] = input_kwargs.pop("stop", None)
 
         if stop is not None:
@@ -177,6 +178,8 @@ def _process_args(
             attention_mask=attention_mask,
             generation_config=GenerationConfig(**generating_args),
         )
+        if seed is not None:
+            gen_kwargs["_seed"] = seed
 
         mm_inputs = template.mm_plugin.get_mm_inputs(**mm_input_dict, batch_ids=[prompt_ids], processor=processor)
         for key, value in mm_inputs.items():
@@ -237,6 +240,10 @@ def _chat(
             audios,
             input_kwargs,
         )
+        seed = gen_kwargs.pop("_seed", None)
+        if seed is not None:
+            set_seed(seed)
+
         generate_output = model.generate(**gen_kwargs)
         if isinstance(generate_output, tuple):
             generate_output = generate_output[1][0]  # post-process the minicpm_o output
@@ -292,6 +299,10 @@ def _stream_chat(
             audios,
             input_kwargs,
         )
+        seed = gen_kwargs.pop("_seed", None)
+        if seed is not None:
+            set_seed(seed)
+
         streamer = TextIteratorStreamer(
             tokenizer,
             skip_prompt=True,

src/llamafactory/chat/sglang_engine.py

@@ -171,6 +171,7 @@ async def _generate(
         skip_special_tokens: Optional[bool] = input_kwargs.pop("skip_special_tokens", None)
         max_length: Optional[int] = input_kwargs.pop("max_length", None)
         max_new_tokens: Optional[int] = input_kwargs.pop("max_new_tokens", None)
+        seed: Optional[int] = input_kwargs.pop("seed", None)
         stop: Optional[Union[str, list[str]]] = input_kwargs.pop("stop", None)
 
         if num_return_sequences != 1:
@@ -205,6 +206,8 @@ async def _generate(
             if skip_special_tokens is not None
             else self.generating_args["skip_special_tokens"],
         }
+        if seed is not None:
+            sampling_params["seed"] = seed
 
         def stream_request():
             json_data = {

src/llamafactory/chat/vllm_engine.py

@@ -144,6 +144,7 @@ async def _generate(
         skip_special_tokens: Optional[bool] = input_kwargs.pop("skip_special_tokens", None)
         max_length: Optional[int] = input_kwargs.pop("max_length", None)
         max_new_tokens: Optional[int] = input_kwargs.pop("max_new_tokens", None)
+        seed: Optional[int] = input_kwargs.pop("seed", None)
         stop: Optional[Union[str, list[str]]] = input_kwargs.pop("stop", None)
 
         if length_penalty is not None:
@@ -163,7 +164,7 @@ async def _generate(
         if max_new_tokens:
             max_tokens = max_new_tokens
 
-        sampling_params = SamplingParams(
+        sampling_kwargs = dict(
             n=num_return_sequences,
             repetition_penalty=(
                 repetition_penalty if repetition_penalty is not None else self.generating_args["repetition_penalty"]
@@ -179,6 +180,10 @@ async def _generate(
             if skip_special_tokens is not None
             else self.generating_args["skip_special_tokens"],
         )
+        if seed is not None:
+            sampling_kwargs["seed"] = seed
+
+        sampling_params = SamplingParams(**sampling_kwargs)
 
         multi_modal_data = {}
         if images is not None:  # add image features

src/llamafactory/webui/chatter.py

@@ -203,6 +203,7 @@ def stream(
         max_new_tokens: int,
         top_p: float,
         temperature: float,
+        seed: float | int,
         skip_special_tokens: bool,
         escape_html: bool,
         enable_thinking: bool,
@@ -225,6 +226,7 @@ def stream(
                 max_new_tokens=max_new_tokens,
                 top_p=top_p,
                 temperature=temperature,
+                seed=None if seed is None or int(seed) < 0 else int(seed),
                 skip_special_tokens=skip_special_tokens,
             ):
                 response += new_text

src/llamafactory/webui/components/chatbot.py

@@ -85,6 +85,7 @@ def create_chat_box(
                 max_new_tokens = gr.Slider(minimum=8, maximum=8192, value=1024, step=1)
                 top_p = gr.Slider(minimum=0.01, maximum=1.0, value=0.7, step=0.01)
                 temperature = gr.Slider(minimum=0.01, maximum=1.5, value=0.95, step=0.01)
+                seed = gr.Number(value=-1, precision=0)
                 skip_special_tokens = gr.Checkbox(value=True)
                 escape_html = gr.Checkbox(value=True)
                 enable_thinking = gr.Checkbox(value=True)
@@ -110,6 +111,7 @@ def create_chat_box(
             max_new_tokens,
             top_p,
             temperature,
+            seed,
             skip_special_tokens,
             escape_html,
             enable_thinking,
@@ -135,6 +137,7 @@ def create_chat_box(
             max_new_tokens=max_new_tokens,
             top_p=top_p,
             temperature=temperature,
+            seed=seed,
             skip_special_tokens=skip_special_tokens,
             escape_html=escape_html,
             enable_thinking=enable_thinking,

src/llamafactory/webui/components/eval.py

@@ -55,10 +55,13 @@ def create_eval_tab(engine: "Engine") -> dict[str, "Component"]:
         max_new_tokens = gr.Slider(minimum=8, maximum=4096, value=512, step=1)
         top_p = gr.Slider(minimum=0.01, maximum=1, value=0.7, step=0.01)
         temperature = gr.Slider(minimum=0.01, maximum=1.5, value=0.95, step=0.01)
+        eval_seed = gr.Textbox(value="42")
         output_dir = gr.Textbox()
 
-    input_elems.update({max_new_tokens, top_p, temperature, output_dir})
-    elem_dict.update(dict(max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature, output_dir=output_dir))
+    input_elems.update({max_new_tokens, top_p, temperature, eval_seed, output_dir})
+    elem_dict.update(
+        dict(max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature, eval_seed=eval_seed, output_dir=output_dir)
+    )
 
     with gr.Row():
         cmd_preview_btn = gr.Button()

src/llamafactory/webui/components/train.py

@@ -52,15 +52,17 @@ def create_train_tab(engine: "Engine") -> dict[str, "Component"]:
         learning_rate = gr.Textbox(value="5e-5")
         num_train_epochs = gr.Textbox(value="3.0")
         max_grad_norm = gr.Textbox(value="1.0")
+        train_seed = gr.Textbox(value="42")
         max_samples = gr.Textbox(value="100000")
         compute_type = gr.Dropdown(choices=["bf16", "fp16", "fp32", "pure_bf16"], value="bf16")
 
-    input_elems.update({learning_rate, num_train_epochs, max_grad_norm, max_samples, compute_type})
+    input_elems.update({learning_rate, num_train_epochs, max_grad_norm, train_seed, max_samples, compute_type})
     elem_dict.update(
         dict(
             learning_rate=learning_rate,
             num_train_epochs=num_train_epochs,
             max_grad_norm=max_grad_norm,
+            train_seed=train_seed,
             max_samples=max_samples,
             compute_type=compute_type,
         )

src/llamafactory/webui/locales.py

@@ -539,6 +539,28 @@
             "info": "勾配クリッピングのためのノルム。",
         },
     },
+    "train_seed": {
+        "en": {
+            "label": "Seed",
+            "info": "Random seed for training.",
+        },
+        "ru": {
+            "label": "Seed",
+            "info": "Random seed for training.",
+        },
+        "zh": {
+            "label": "随机种子",
+            "info": "训练使用的随机种子。",
+        },
+        "ko": {
+            "label": "Seed",
+            "info": "Random seed for training.",
+        },
+        "ja": {
+            "label": "Seed",
+            "info": "Random seed for training.",
+        },
+    },
     "max_samples": {
         "en": {
             "label": "Max samples",
@@ -2708,6 +2730,45 @@
             "label": "温度",
         },
     },
+    "seed": {
+        "en": {
+            "label": "Generation seed (-1 for random)",
+        },
+        "ru": {
+            "label": "Generation seed (-1 = random)",
+        },
+        "zh": {
+            "label": "生成随机种子（-1 表示随机）",
+        },
+        "ko": {
+            "label": "Generation seed (-1 = random)",
+        },
+        "ja": {
+            "label": "Generation seed (-1 = random)",
+        },
+    },
+    "eval_seed": {
+        "en": {
+            "label": "Seed",
+            "info": "Random seed for evaluation and prediction.",
+        },
+        "ru": {
+            "label": "Seed",
+            "info": "Random seed for evaluation and prediction.",
+        },
+        "zh": {
+            "label": "随机种子",
+            "info": "评估和预测使用的随机种子。",
+        },
+        "ko": {
+            "label": "Seed",
+            "info": "Random seed for evaluation and prediction.",
+        },
+        "ja": {
+            "label": "Seed",
+            "info": "Random seed for evaluation and prediction.",
+        },
+    },
     "skip_special_tokens": {
         "en": {
             "label": "Skip special tokens",

src/llamafactory/webui/runner.py

@@ -51,6 +51,13 @@
     from .manager import Manager
 
 
+def _parse_seed(value: Any, default: int = 42) -> int:
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return default
+
+
 class Runner:
     r"""A class to manage the running status of the trainers."""
 
@@ -147,6 +154,7 @@ def _parse_train_args(self, data: dict["Component", Any]) -> dict[str, Any]:
             learning_rate=float(get("train.learning_rate")),
             num_train_epochs=float(get("train.num_train_epochs")),
             max_samples=int(get("train.max_samples")),
+            seed=_parse_seed(get("train.train_seed")),
             per_device_train_batch_size=get("train.batch_size"),
             gradient_accumulation_steps=get("train.gradient_accumulation_steps"),
             lr_scheduler_type=get("train.lr_scheduler_type"),
@@ -316,6 +324,7 @@ def _parse_eval_args(self, data: dict["Component", Any]) -> dict[str, Any]:
             max_new_tokens=get("eval.max_new_tokens"),
             top_p=get("eval.top_p"),
             temperature=get("eval.temperature"),
+            seed=_parse_seed(get("eval.eval_seed")),
             output_dir=get_save_dir(model_name, finetuning_type, get("eval.output_dir")),
             trust_remote_code=True,
             ddp_timeout=180000000,