Set ccl and KMP param in simple launch

[jiqing-feng]

Don't know why it assigns CCL_WORKER_COUNT only when machine> 1 because 1 CPU machine can also run distributed training or Tensor Parallelism.

I also added KMP params to get better performance on CPU.

With this PR we can run transformers TP model and got 40% speed-up on Intel 4th Gen Xeon.
The accelerate config is :

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: MULTI_CPU
downcast_bf16: 'no'
enable_cpu_affinity: false
ipex_config:
  ipex: false
machine_rank: 0
main_training_function: main
mixed_precision: 'no'
mpirun_config:
  mpirun_ccl: '1'
  mpirun_hostfile: /home/jiqingfe/hostfile
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: true

The script is as follows
accelerate launch script.py

import os
import torch.distributed as dist
from transformers import AutoTokenizer, AutoModelForCausalLM
import oneccl_bindings_for_pytorch

import time
import torch

print(f"Using {torch.get_num_threads()} threads (PyTorch)")
print(f"OMP_NUM_THREADS={os.getenv('OMP_NUM_THREADS')}")

model_id = "meta-llama/Llama-3.1-8B-Instruct"

os.environ['RANK'] = str(os.environ.get('PMI_RANK', 0))
os.environ['WORLD_SIZE'] = str(os.environ.get('PMI_SIZE', 1))

def main(is_tp, rank, world_size) -> None:
    print("is_tp, rank, world_size: ", is_tp, rank, world_size)
    model_kwargs = dict(torch_dtype=torch.bfloat16)
    if is_tp:
        model_kwargs["tp_plan"] = "auto"
    else:
        model_kwargs["device_map"] = "cpu"

    # Retrieve tensor parallel model
    model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
    if dist.is_initialized():
        print("Backend:", dist.get_backend())
    else:
        print("Distributed process group is not initialized.")

    # Prepare input tokens
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    prompt = "It is done, and submitted. You can play 'Survival of the Tastiest' on Android, and on the web. Playing on the web works, but you have to simulate multiple touch for table moving and that can be a bit confusing. There is a lot I'd like to talk about. I will go through every topic, insted of making the typical what went right/wrong list. Concept Working over the theme was probably one of the hardest tasks which I had to face. Originally, I had an idea of what kind of game I wanted to develop, gameplay wise - something with a lot of enemies/actors"
    inputs = tokenizer(prompt, return_tensors="pt", max_length=512).to(model.device)
    print(f"inpu shape is {inputs.input_ids.shape}")

    model.generation_config.cache_implementation = "static"

    if is_tp:
        model.config.hidden_size = model.config.hidden_size // world_size
        model.config.num_key_value_heads = model.config.num_key_value_heads // world_size

    for i in range(1):
        with torch.no_grad():
            start = time.time()
            outputs = model.generate(**inputs, do_sample=False, max_new_tokens=128, min_new_tokens=128)
            end = time.time()
            print(f"time cost {(end-start)*1000} ms")

    # warm-up
    if is_tp:
        dist.barrier()

    if rank == 0:
        print(tokenizer.batch_decode(outputs, skip_special_tokens=True))

    model.forward = torch.compile(model.forward)
    # warm-up
    if is_tp:
        dist.barrier()

    for i in range(4):
        with torch.no_grad():
            start = time.time()
            outputs = model.generate(**inputs, do_sample=False, max_new_tokens=128, min_new_tokens=128)
            if is_tp:
                dist.barrier()

            end = time.time()
            print(f"time cost {(end-start)*1000} ms")

    if rank == 0:
        print(tokenizer.batch_decode(outputs, skip_special_tokens=True))


if __name__ == "__main__":
    rank = int(os.environ["RANK"]) if "RANK" in os.environ else 0
    world_size = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
    is_tp = world_size > 1
    main(is_tp, rank, world_size)

[jiqing-feng]

@sywangyi @yao-matrix . Please review this PR, thanks!

[yao-matrix]

1. does ipex: true work？
2. in the case we don't have accelerate config, what's the behavior?

[jiqing-feng]

Hi @yao-matrix . I have verified that both ipex is True and no config can work as before.

[jiqing-feng]

Hi @SunMarc . Could you please review this PR? Thanks!

[SunMarc]

Thanks, left a couple of comments

[SunMarc]

LGTM, just a small nit

[jiqing-feng]

LGTM, just a small nit

Fixed!

[SunMarc]

Thanks ! LGTM !

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.