[PZ COMPETITION] code submission by losercheems

[LoserCheems]

Hi, the code has been submitted through this PR.
The checkpoints and logs have been published on HuggingFace.

• perliminary
• final

[gemini-code-assist]

Summary of Changes

Hello @LoserCheems, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request establishes a comprehensive training infrastructure for the OpenSeek model, covering both its initial continue pre-training and subsequent KTO-style preference alignment. It provides all necessary configuration files, data preparation scripts, and training entry points, designed for efficient execution on multi-GPU environments using accelerate and DeepSpeed. The changes enable reproducible model development from raw data to a fine-tuned checkpoint, complete with evaluation examples and merging capabilities.

Highlights

• OpenSeek KTO Alignment Pipeline: Introduced a complete Kahneman–Tversky Optimization (KTO) alignment fine-tuning pipeline for the OpenSeek model, including scripts for asset acquisition, dataset transformation, and alignment training.
• OpenSeek Continue Pre-training Workflow: Added a three-stage continue pre-training workflow for the OpenSeek model, encompassing asset download, mixed dataset building, distributed training, and model checkpoint merging (model soup).
• Distributed Training Configurations: Provided accelerate configuration files for both DeepSpeed ZeRO-2 (for KTO) and DDP (for pre-training) to enable efficient multi-GPU training.
• Data Preparation Scripts: Developed dedicated Python scripts for downloading base models and datasets, as well as transforming raw datasets into formats suitable for KTO preference alignment and pre-training with mixed ratios and sequence packing.
• SLURM Integration: Included train.sh scripts for both pipelines to facilitate job submission on SLURM clusters, setting up necessary environment variables and launching accelerate commands.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[Copilot]

Pull Request Overview

This PR implements a comprehensive training and alignment pipeline for the OpenSeek model, adding configuration files, data processing scripts, and training components for both pretraining and preference-based fine-tuning (KTO).

• Creates complete training infrastructure with distributed training configurations and dataset processing utilities
• Implements both preliminary pretraining pipeline and final KTO alignment workflow
• Adds comprehensive documentation and example configurations for model training and evaluation

Reviewed Changes

Copilot reviewed 16 out of 18 changed files in this pull request and generated 10 comments.

Show a summary per file

File	Description
openseek/competition/pz/losercheems/preliminary/utils/training_args_configs.py	Defines PTConfig class for pretraining configuration parameters
openseek/competition/pz/losercheems/preliminary/trainer/pt_dpsk.py	Main pretraining script with Trainer setup and dataset processing
openseek/competition/pz/losercheems/preliminary/train.sh	SLURM batch script for distributed training launch
openseek/competition/pz/losercheems/preliminary/scripts/merge.py	Model checkpoint merging utility with multiple strategies
openseek/competition/pz/losercheems/preliminary/scripts/download.py	Dataset downloading and preprocessing script
openseek/competition/pz/losercheems/preliminary/recipes/openseek/config.yaml	Training configuration with model and hyperparameter settings
openseek/competition/pz/losercheems/preliminary/recipes/accelerate_configs/ddp.yaml	Distributed training configuration for Accelerate
openseek/competition/pz/losercheems/preliminary/processor/pt_datasets_process.py	Dataset processing utilities for mixing and tokenization
openseek/competition/pz/losercheems/preliminary/README.md	Technical documentation for the pretraining workflow
openseek/competition/pz/losercheems/final/trainer/kto.py	KTO alignment training script
openseek/competition/pz/losercheems/final/train.sh	Training launch script for KTO alignment
openseek/competition/pz/losercheems/final/scripts/kto_datasets_process.py	Dataset processing for KTO preference format
openseek/competition/pz/losercheems/final/scripts/download.py	Asset download script for KTO training
openseek/competition/pz/losercheems/final/recipes/openseek/config.yaml	KTO training configuration
openseek/competition/pz/losercheems/final/recipes/accelerate_configs/zero2.yaml	DeepSpeed ZeRO-2 configuration
openseek/competition/pz/losercheems/final/README.md	Technical documentation for KTO alignment workflow

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[Copilot]

This import will fail because 'utils' is not in the Python path. The import should be a relative import: from ..utils.training_args_configs import PTConfig or an absolute import from the package root.

Suggested change

	from utils.training_args_configs import PTConfig
	from ..utils.training_args_configs import PTConfig

[Copilot]

This import references 'small_doge.processor' but should import from the local processor module: from ..processor.pt_datasets_process import mix_datasets_by_ratio as mix_pt_datasets.

Suggested change

	from small_doge.processor import mix_pt_datasets
	from ..processor.pt_datasets_process import mix_datasets_by_ratio as mix_pt_datasets

[Copilot]

This commented-out code should be removed as it clutters the codebase and provides no value in its current state.

Suggested change

# if model_args.model_name_or_path is not None and model_args.model_name_or_path.endswith("checkpoint") else causal_lm_class(config=config).to(torch_dtype)

[Copilot]

The time limit '114514' appears to be a placeholder or meme number rather than a realistic job duration. This should be set to an appropriate time limit for the training job.

Suggested change

	#SBATCH -t 114514
	#SBATCH -t 1-00:00:00

[Copilot]

The checkpoint-5000 path is duplicated in the list. This will cause the same checkpoint to be loaded twice during merging, which is likely unintentional.

Suggested change

checkpoint_paths.append(r"./data/OpenSeek-1.4B-A0.4B/checkpoint-5000")

[Copilot]

This import path is incorrect. It should be a relative import: from ..processor.pt_datasets_process import mix_datasets_by_ratio as mix_pt_datasets or an absolute import from the package root.

Suggested change

	from preliminary.processor.pt_datasets_process import mix_datasets_by_ratio as mix_pt_datasets
	from ..processor.pt_datasets_process import mix_datasets_by_ratio as mix_pt_datasets

[Copilot]

The calculate_total_ratio function expects a dictionary but datasets_and_ratios is a list of dictionaries. This line should be: total_ratio = sum(list(dataset.values())[0] for dataset in datasets_and_ratios).

Suggested change

	def calculate_total_ratio(datasets_and_ratios):
	return sum(item for item in datasets_and_ratios.values())
	total_ratio = sum(calculate_total_ratio(dataset) for dataset in datasets_and_ratios)
	total_ratio = sum(list(dataset.values())[0] for dataset in datasets_and_ratios)

[Copilot]

The number of processes (8) doesn't match the GPU allocation in train.sh (4 GPUs). These values should be consistent - either both should be 4 or both should be 8.

Suggested change

	num_processes: 8
	num_processes: 4

[Copilot]

This line assumes the dataset already has an 'input_ids' column, but the dataset may not be preprocessed yet. This should be moved after the prepare_dataset call or made conditional on whether the dataset is already processed.

[Copilot]

This import is unused in the code and references an external module that may not exist. The import should be removed since the dataset is loaded directly with datasets.load_from_disk.

Suggested change

from small_doge.processor import mix_dpo_datasets

[gemini-code-assist]

Code Review

This pull request introduces a comprehensive set of files for pre-training and KTO alignment of the OpenSeek model. The changes include scripts for data downloading and processing, configuration files for training, and the training scripts themselves, along with detailed documentation. The overall structure is good, but there are several areas for improvement, particularly concerning hardcoded paths, code clarity, and reproducibility. I've identified several critical and high-severity issues related to hardcoded absolute paths in configuration and scripts, which will prevent others from running this code without modification. There are also some logical errors, such as a duplicated checkpoint in the merge script and confusing data processing logic. I've provided specific suggestions to address these points to make the pipeline more robust and portable.

[gemini-code-assist]

The dataset path /workspace/datasets/OpenSeek-Pretrain-30B is hardcoded as an absolute path. This severely impacts reproducibility, as the training will fail on any machine where this exact path does not exist. Please use a relative path or an environment variable.

  - ./datasets/OpenSeek-Pretrain-30B: 1.0

[gemini-code-assist]

The script contains a hardcoded absolute path /root/code/small-doge/datasets/AI-MO/NuminaMath-CoT for loading the dataset and a hardcoded relative path for saving. This severely impacts portability and reproducibility. These paths should be parameterized using command-line arguments (e.g., with argparse) to make the script reusable.

Suggested change

	numina_math_cot = load_from_disk("/root/code/small-doge/datasets/AI-MO/NuminaMath-CoT")
	print(numina_math_cot)
	numina_math_cot = numina_math_cot.map(process, num_proc=4).select_columns(["prompt", "completion", "label"])
	print(numina_math_cot)
	print(numina_math_cot["train"][0])
	numina_math_cot.save_to_disk("./datasets/AI-MO/NuminaMath-CoT-preference")
	import argparse
	if __name__ == "__main__":
	parser = argparse.ArgumentParser(description="Process dataset for KTO training.")
	parser.add_argument("--input_path", type=str, required=True, help="Path to the input dataset.")
	parser.add_argument("--output_path", type=str, required=True, help="Path to save the processed dataset.")
	args = parser.parse_args()
	numina_math_cot = load_from_disk(args.input_path)
	print(numina_math_cot)
	numina_math_cot = numina_math_cot.map(process, num_proc=4).select_columns(["prompt", "completion", "label"])
	print(numina_math_cot)
	print(numina_math_cot["train"][0])
	numina_math_cot.save_to_disk(args.output_path)

[gemini-code-assist]

The path to checkpoint-5000 is duplicated in the checkpoint_paths list. This is likely a copy-paste error and will cause this checkpoint's weights to be counted twice, skewing the average during the merge process.

[gemini-code-assist]

This line prematurely filters the dataset to only the input_ids column before calling prepare_dataset. This prevents the tokenization logic inside prepare_dataset (which expects a text field) from ever running, making the code confusing and brittle if a raw text dataset is ever used. This selection should happen inside prepare_dataset after tokenization and before packing.

[gemini-code-assist]

The KTOTrainer from trl does not accept a processing_class argument. It seems you intended to pass the tokenizer, but the argument name is incorrect and it will be ignored. The tokenizer is usually handled through the DataCollator or is implicitly used by the trainer if needed. Since KTOTrainer doesn't require it in its constructor, this line can be removed.

[gemini-code-assist]

The indentation for final_result.json in the directory overview seems incorrect. It appears to be a file inside the <benchmark_name> directory, but based on the description, it should likely be at the same level as <benchmark_name>.

Suggested change

	final_result.json # Aggregated metrics summary
	<benchmark_name>/ # Per-benchmark JSONL + metrics
	final_result.json # Aggregated metrics summary
	<benchmark_name>/ # Per-benchmark JSONL + metrics

[gemini-code-assist]

This large block of commented-out datasets_and_ratios makes the configuration file noisy and difficult to read. If this is for reference or an example, it should be moved to the documentation. Otherwise, it should be removed to improve clarity.

[gemini-code-assist]

The docstring and comments are in Chinese, which is inconsistent with the rest of the project (e.g., READMEs, other scripts) which is in English. For better maintainability and accessibility for all contributors, please translate these to English.

Suggested change

	"""
	合并多个checkpoints
	Args:
	checkpoint_paths: 要合并的checkpoint路径列表
	output_path: 合并后模型的保存路径
	merge_method: 合并方法，支持 "average", "last", "weighted_average"
	"""
	"""
	Merges multiple checkpoints.
	Args:
	checkpoint_paths: A list of checkpoint paths to be merged.
	output_path: The path to save the merged model.
	merge_method: The method for merging, supports "average", "last", "weighted_average".
	"""

[gemini-code-assist]

The SBATCH time limit -t 114514 appears to be a placeholder or a meme value. For professionalism and proper resource management on a SLURM cluster, this should be set to a realistic time estimate for the job, e.g., 24:00:00 for 24 hours.

Suggested change

	#SBATCH -t 114514
	#SBATCH -t 24:00:00

[gemini-code-assist]

Using yaml.load with FullLoader is unsafe as it can execute arbitrary code. For loading simple configuration files, yaml.safe_load should be used instead to prevent potential security vulnerabilities.

Suggested change

	model_config = yaml.load(
	open(config_path, "r", encoding="utf-8"), Loader=yaml.FullLoader
	)["model_config"]
	model_config = yaml.safe_load(
	open(config_path, "r", encoding="utf-8"),
	)["model_config"]