developer_readme.md

Welcome

Thank you for contributing to MoCo! This is the guide for contriubting your model collaboration algorithm to MoCo. If you are looking to contribute a dataset/evaluation, please check out docs/eval_readme.md.

Where to put your code

You will move to your own branch via git, create a new script <type>_<method_name>.py in the model_collaboration/method/ directory, and implement your method there. <type> should be in api, text, logit, weight, or very occassionally mixed (explain why in the pull request if you believe your approach is mixed), <method_name> should be a one-or-two-word phrase.

After implementing and testing your method, before the pull request, add a description of your method to docs/user_readme.md and follow the format there.

You may only submit one method per pull request. You may only modify the .py of your approach and the docs/user_readme.md. If you have suggestions for other methods / shared helper functions, please open an issue.

A template

A template of an approach is available in the model_collaboration/method/sample_approach.py file. You can use this as a starting point for your own method.

Helper functions

We provide a lot of helper functions for you to build on: distributed generation with multiple models and multiple GPUs, distributed sft & dpo, Particle Swarm Optimization for optimizing a bunch of weights/hyperparameters, etc.

method/distributed_generation.py

Import: from method import distributed_generation Usage: distributed_generation.distributed_generation(list_of_model_name, list_of_input_list, list_of_gpu_id, max_response_length=None)

• list_of_model_name: a list of model names (strings) that are compatible with transformers.AutoModelForCausalLM.from_pretrained(). Should be model_names that users provide in config files.
• list_of_input_list: a list of list of strings, size len(model_names) * any. Each inner list contains the inputs to be generated by the corresponding model.
• list_of_gpu_id: a list of GPU ids (integers) to use for generation. The length of this list could be anything: smaller, equal, larger than len(model_names). Larger is a waste. Smaller means models are assigned to GPUs in a round-robin fashion. Equal is perhaps best for parallelism. Users provide gpu_ids in config files, you should pass in it or a subset of it.
• max_response_length: HIGHLY optional, you probably don't need to set this. If set, it overrides the default max response length provided in the config. Only when it makes sense to generate another amount of tokens, shorter/longer than the response length the user asked for, to set this. Should be rare.
• It can work with one model! Just do [model_name], [input_list], [gpu_id] in that case.

data/eval.py

Import: from data import eval

Taking the evaluation inputs as a list: input_list = eval.prepare_inputs(task, task_type, split, ratio=1.0)

• task: string, the name of the evaluation task, e.g., agieval
• task_type: string, the type of the evaluation task, e.g., multiple_choice. See data/eval_readme.md for valid (task, task_type) combinations.
• split: string, the data split to use, dev or test.
• ratio: You almost certainly don't need to use this. Float between 0 and 1.0, the fraction of data to use. Default 1.0, use all data.
• Returns: input_list, a list of strings.

Evaluating a list of outputs to those inputs: scores = eval.get_scores(task, task_type, split, outputs, ratio=1.0)

• task, task_type, split, ratio: same as above.
• outputs: a list of strings, the model-generated outputs corresponding to the inputs obtained from eval.prepare_inputs().
• Returns: scores, a list of floats, the evaluation scores for each input-output pair.

utils/distributed_sft.py

Import: from utils import distributed_sft Usage: distributed_sft.distributed_sft(list_of_model_names, list_of_sft_data_paths, list_of_gpu_ids, list_of_output_model_paths, batch_size=1, gradient_accumulation_steps=16, learning_rate=1e-5, epoch=3)

• list_of_model_names: a list of model names (strings) that are compatible with transformers.AutoModelForCausalLM.from_pretrained(), that you are going to fine-tune.
• list_of_sft_data_paths: a list of paths to SFT data files (strings), each corresponding to a model in list_of_model_names. The data files should be JSONL with each line {"prompt":..., "completion":...}.
• list_of_gpu_ids: a list of GPU ids (integers) to use for SFT. The length of this list could be anything: smaller, equal, larger than len(model_names). Larger is a waste. Smaller means models are assigned to GPUs in a round-robin fashion. Equal is perhaps best for parallelism. Users provide gpu_ids in config files, you should pass in it or a subset of it.
• list_of_output_model_paths: a list of paths (strings) to save the fine-tuned models, each corresponding to a model in list_of_model_names. Please save them in model_collaboration/logs/<your_method_name>/ folder to avoid conflict.
• batch_size, gradient_accumulation_steps, learning_rate, epoch: standard SFT hyperparameters, you can set them as you like. batch_size * gradient_accumulation_steps is the effective batch size.

utils/distributed_dpo.py

Import: from utils import distributed_dpo Usage: distributed_dpo.distributed_dpo(list_of_model_names, list_of_dpo_data_paths, list_of_gpu_ids, list_of_output_model_paths, batch_size=1, gradient_accumulation_steps=16, learning_rate=1e-6, epoch=1)

• list_of_model_names: a list of model names (strings) that are compatible with transformers.AutoModelForCausalLM.from_pretrained(), that you are going to fine-tune with DPO. Can be HuggingFace Hub identifiers or local paths. If a model path is a LoRA adapter, it will be automatically merged into the base model before training a new adapter.
• list_of_dpo_data_paths: a list of paths to DPO data files (strings), each corresponding to a model in list_of_model_names. The data files should be JSONL with each line {"prompt":..., "chosen":..., "rejected":...}. Alternatively, if the data uses "instruction" as the prompt field, it will be automatically renamed to "prompt" for compatibility.
• list_of_gpu_ids: a list of GPU ids (integers) to use for DPO. The length of this list could be anything: smaller, equal, larger than len(model_names). Larger is a waste. Smaller means models are assigned to GPUs in a round-robin fashion. Equal is perhaps best for parallelism. Users provide gpu_ids in config files, you should pass in it or a subset of it.
• list_of_output_model_paths: a list of paths (strings) to save the DPO-trained models (as LoRA adapters), each corresponding to a model in list_of_model_names. Please save them in model_collaboration/logs/<your_method_name>/ folder to avoid conflict.
• batch_size, gradient_accumulation_steps, learning_rate, epoch: standard DPO hyperparameters, you can set them as you like. batch_size * gradient_accumulation_steps is the effective batch size. The default learning rate is 1e-6 (typically lower than SFT).

utils/numeric_swarm.py

Use case: Particle Swarm Optimization (PSO) for optimizing a set of numeric parameters. For example in model merging x = \sum_{i=1}^n w_i * x_i, you may want to optimize the weights w_i using PSO.

Import: from utils.numeric_swarm import NumericSwarm Creating a PSO optimizer: swarm = NumericSwarm(dimension, population, starting_velocity_mode="random", weight_randomness=True, inertia=0.2, cognitive_coeff=0.3, social_coeff=0.4, repel_coeff=0.05, step_length=0.5, repel_term=True, step_length_factor=0.95, minimum_step_length=0.1, patience=5, restart_patience=3)

• dimension: integer, the number of parameters to optimize, n in the use case.
• population: integer, the number of particles in the swarm. Essentially how many candidate solutions to maintain in each iteration. 5-10 recommended, larger is more thorough but much slower.
• for all other hyperparameters, a default value is provided. Learn more (if you want to) at link.

Getting the current positions of all particles: positions = swarm.get_particles()

• positions will be a list of lists of floats, size population * dimension. Each inner list is the current position (i.e., the current candidate solution) of a particle.

Updating the swarm with new scores for each particle: terminate_signal = swarm.update(scores)

• scores: a list of floats, the scores (the higher the better) for each particle in the swarm. Length should be equal to population.
• terminate_signal: boolean, whether the optimization should terminate because of patience, for early stopping.

Getting the global best position found so far: best_position = swarm.get_global_best_particle()

• best_position: a list of floats, size dimension, the best candidate solution found so far.

There is no explicit max iterations: you just call swarm.update() for that many of times and stop when you want and use swarm.get_global_best_particle() to get the best solution found in the end.

utils/logit_arithmetic.py

Use case: you are fusing/contrasting the next-token logits from multiple models, e.g., logits = w_1 * logits_1 + w_2 * logits_2 - w_3 * logits_3, and genearte text from that joint distribution. Other logit operations (e.g. taking a confidence threshold over logit to decide something) is not supported here, implement in your own script.

Import: from utils import logit_arithmetic Creating a logit arithmetic object: logit_operator = logit_arithmetic.LogitArithmetic(model_names, model_devices, tokenizer)

• model_names: a list of model names (strings) that are compatible with transformers.AutoModelForCausalLM.from_pretrained(), that you are going to use for logit arithmetic.
• model_devices: a list of device strings (e.g., ['cuda:0', 'cuda:1', ...]) to load each model onto. Length should be equal to model_names.
• tokenizer: a tokenizer object compatible with the models, e.g., transformers.AutoTokenizer.from_pretrained(model_name).

Generating responses with logit arithmetic: outputs = logit_operator.batch_generate(prompts, tokenizer, batch_size=1, max_new_tokens=100, do_sample=True, temperature=1.0, arithmetic_func=average_logits)

• prompts: a list of strings, the input prompts to generate from.
• tokenizer: the tokenizer object used to encode the prompts and decode the outputs.
• batch_size, max_new_tokens, do_sample, temperature: standard generation parameters.
• arithmetic_func: a function that takes in a list of logits tensors (one per model) and returns a single logits tensor after performing the desired arithmetic operation. An example function average_logits is provided in logit_arithmetic.py that simply averages the logits from all models, and is the default. See how a custom logit arithmetic function can be defined in method/logit_logit_contrastive.py for reference.

If your method would like to use some custom helper functions, please put them under model_collaboration/utils/.