Will Trainer.predict() return data in the same order as the original dataset during multi-machine and multi-gpus inference?
Author: deepwebney
Pull request: no
Labels: name: trainer

name: Feature request

name: Accelerate

### Feature request

I want to use accelerate for multi-machine and multi-gpus inference. Since using trainer.predict does not return the original inference data, only the inference results, I am not sure if their order can still be maintained in the case of multi-machine and multi-gpus.

### Motivation

I want to use accelerate for multi-machine and multi-gpus inference. Since using trainer.predict does not return the original inference data, only the inference results, I am not sure if their order can still be maintained in the case of multi-machine and multi-gpus.

### Your contribution

this is my code：
def main(args):
    def tokenize_function(examples):
        return tokenizer(examples['text'], padding=False, truncation=False)

    # Step 1: Load the model
    model = LlamaForSequenceClassification.from_pretrained(args.model_name, num_labels=args.num_class, problem_type='multi_label_classification', trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(args.model_name, use_fast=True)

    # Step 2: Load the dataset

    eval_dataset = load_dataset(args.eval_dataset_type, data_files=args.eval_dataset_name)

    tokenized_dataset = eval_dataset.map(tokenize_function, batched=True, batch_size=10000)


    # Step 3: Define the training arguments
    training_args = MyTrainingArguments(
        hdfs_path=args.hdfs_path,
        output_dir=args.output_dir,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        logging_steps=args.logging_steps,
        num_train_epochs=args.num_train_epochs,
        max_steps=args.max_steps,
        report_to=args.report_to,
        save_steps=args.save_steps,
        save_total_limit=args.save_total_limit,
        gradient_checkpointing=args.gradient_checkpointing,
        logging_dir=args.logging_dir,
        bf16=True,
        evaluation_strategy="steps",
        eval_steps=args.eval_steps 
    )
    
    # Step 5: Define the trainer 
    trainer = Trainer(model=model,
        args=training_args,
        eval_dataset=tokenized_dataset['train'],
        compute_metrics=compute_metrics,
        data_collator=DataCollatorWithPadding(tokenizer=tokenizer, max_length=args.max_length),
        tokenizer=tokenizer,)

    model, trainer = accelerator.prepare(model, trainer)
    thresholds = [0.177, 0.041, 0.299, 0.023, 0.049, 0.012, 0.057, 0.212, 0.187, 0.044, 0.107, 0.035, 0.257, 0.19, 0.258, 0.26, 0.166, 0.097, 0.263, 1.0, 0.549, 0.22, 0.03, 0.294, 0.232, 0.524, 0.113, 0.028, 0.064, 0.135, 0.289, 0.121, 0.016, 0.32, 0.095]
    thresholds = np.array(thresholds)

    with torch.no_grad():
        predictions = trainer.predict(tokenized_dataset['train'])
        score = predictions.predictions
        score = 1 / (1 + np.exp(-score))
        pred = np.where(score >= thresholds, 1, 0)
        pred = pred.tolist()

Comments:
cc @muellerzr @SunMarc 

I have the same question. Intuitively, I believe the answer is yes, but I’d still like a confirmation that the order is maintained in a multi-machine or multi-GPU setup.

Hi @muellerzr @SunMarc, I have the same question, would be great to get confirmation on this since I believe this is critical for running inference at scale. Or is there a better approach for merging predictions with the original data?

@PhantomKidding @deepwebney Were either of you able to figure out the answer?

Thanks!

Hi @deepwebney, sorry for the very slow reply, we're catching up on the trainer backlog.

Short answer: yes, `Trainer.predict()` preserves the original dataset order, including on multi-GPU / multi-machine runs. The eval dataloader is never shuffled, see https://github.com/huggingface/transformers/blob/150eb7c9ed4091294c829fa0e9466b090cb0f87f/src/transformers/trainer.py#L1071-L1074
and predictions are gathered per step with `accelerate.gather_for_metrics`  which re-interleaves each rank's shard back into sequential order and truncates the extra rows accelerate adds to even out the last batch.

One note on your script: drop the `accelerator.prepare(model, trainer)` call: the Trainer builds and manages its own `Accelerator` internally, so just launch with `accelerate launch`/`torchrun` and call `trainer.predict()` directly.