[rank0]: ValueError: Your setup doesn't support bf16/gpu.
Author: kadirnar
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.54.0
- Platform: Linux-5.15.0-144-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.34.1
- Safetensors version: 0.5.3
- Accelerate version: 1.9.0
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: MULTI_GPU
        - mixed_precision: bf16
        - use_cpu: False
        - debug: False
        - num_processes: 8
        - machine_rank: 0
        - num_machines: 1
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - enable_cpu_affinity: True
        - fsdp_config: {'fsdp_auto_wrap_policy': 'TRANSFORMER_BASED_WRAP', 'fsdp_backward_prefetch_policy': 'BACKWARD_PRE', 'fsdp_forward_prefetch': True, 'fsdp_offload_params': False, 'fsdp_sharding_strategy': 1, 'fsdp_state_dict_type': 'FULL_STATE_DICT', 'fsdp_transformer_layer_cls_to_wrap': ['LlamaDecoderLayer']}
        - downcast_bf16: False
        - tpu_use_cluster: False
        - tpu_use_sudo: False
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.7.1+cu126 (NA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@eustlb @ArthurZucker @zach-huggingface @SunMarc 

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

code:
```python
import torch
from datasets import load_dataset
from transformers import Trainer, TrainingArguments, AutoTokenizer
import numpy as np
from torch.distributed.fsdp.fully_sharded_data_parallel import FullStateDictConfig
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP, FullStateDictConfig, StateDictType)
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
import yaml
import wandb
from huggingface_hub import HfApi
from accelerate import Accelerator
import math
from liger_kernel.transformers import AutoLigerKernelForCausalLM

config_file = "/mnt/kadirnar/...x/.../configs/llama3_config.yaml"

with open(config_file, "r") as file:
    config = yaml.safe_load(file)

dsn1 = config["text_QA_dataset"]
dsn2 = config["TTS_dataset"]

model_name = config["model_name"]
tokenizer_name = config["tokenizer_name"]

run_name = config["run_name"]
project_name = config["project_name"]
base_repo_id = config["save_folder"]

epochs = config["epochs"]
batch_size = config["batch_size"]
save_steps = config["save_steps"]
pad_token = config["pad_token"]
number_processes = config["number_processes"]
learning_rate = config["learning_rate"]

# Parse ratio from config (e.g., "2:1" -> 2)
ratio_str = config["ratio"]
initial_ratio = int(ratio_str.split(":")[0])
final_ratio = 1  # Target ratio is 1:1


class GradualRatioDataset(Dataset):
    def __init__(self, dataset1, dataset2, batch_total, initial_ratio=2, final_ratio=1, total_steps=None):
        self.dataset1 = dataset1
        self.dataset2 = dataset2
        self.batch_total = batch_total
        self.initial_ratio = initial_ratio
        self.final_ratio = final_ratio
        self.total_steps = total_steps
        
        # Calculate length based on the maximum ratio to ensure we have enough data
        max_ratio = max(initial_ratio, final_ratio)
        num_cycles_ds1 = len(dataset1) // (batch_total * max_ratio)
        num_cycles_ds2 = len(dataset2) // batch_total
        self.num_cycles = min(num_cycles_ds1, num_cycles_ds2)
        
        # Use initial ratio for length calculation
        self.length = self.num_cycles * (initial_ratio + 1) * batch_total
        
        # For tracking current step
        self.current_step = 0

    def set_current_step(self, step):
        """Called by trainer to update current step for ratio calculation"""
        self.current_step = step

    def get_current_ratio(self):
        """Calculate current ratio based on training progress"""
        if self.total_steps is None or self.total_steps == 0:
            return self.initial_ratio
        
        # Linear interpolation from initial_ratio to final_ratio
        progress = min(self.current_step / self.total_steps, 1.0)
        current_ratio = self.initial_ratio - (self.initial_ratio - self.final_ratio) * progress
        return max(int(round(current_ratio)), self.final_ratio)

    def __len__(self):
        return int(self.length)

    def __getitem__(self, index):
        current_ratio = self.get_current_ratio()
        
        # Compute the cycle length in terms of samples with current ratio
        cycle_length = (current_ratio + 1) * self.batch_total
        cycle = index // cycle_length
        pos_in_cycle = index % cycle_length

        if pos_in_cycle < current_ratio * self.batch_total:
            # Text dataset (dataset1)
            batch_in_cycle = pos_in_cycle // self.batch_total
            sample_in_batch = pos_in_cycle % self.batch_total
            ds1_index = cycle * current_ratio * self.batch_total + batch_in_cycle * self.batch_total + sample_in_batch
            
            # Handle index overflow by wrapping around
            if ds1_index >= len(self.dataset1):
                ds1_index = ds1_index % len(self.dataset1)
            
            return self.dataset1[ds1_index]
        else:
            # TTS dataset (dataset2)
            sample_in_batch = pos_in_cycle - current_ratio * self.batch_total
            ds2_index = cycle * self.batch_total + sample_in_batch
            
            # Handle index overflow by wrapping around
            if ds2_index >= len(self.dataset2):
                ds2_index = ds2_index % len(self.dataset2)
            
            return self.dataset2[ds2_index]


class AlternatingDistributedSampler(DistributedSampler):
    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=False):
        super().__init__(dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle)
        self.shuffle = shuffle

    def __iter__(self):
        indices = list(range(len(self.dataset)))
        indices = indices[self.rank:self.total_size:self.num_replicas]
        return iter(indices)


class FSDPTrainer(Trainer):
    def __init__(self, *args, initial_ratio=2, final_ratio=1, **kwargs):
        super().__init__(*args, **kwargs)
        self.repo_id = base_repo_id
        self.api = HfApi()

        self.initial_ratio = initial_ratio
        self.final_ratio = final_ratio
        self.text_step = 0
        self.audio_step = 0
        
        # Calculate total steps for gradual ratio adjustment
        self.total_steps = self.calculate_total_steps()

    def calculate_total_steps(self):
        """Calculate total training steps"""
        num_update_steps_per_epoch = len(self.train_dataset) // (
            self.args.per_device_train_batch_size * self.args.gradient_accumulation_steps * self.args.world_size
        )
        return int(num_update_steps_per_epoch * self.args.num_train_epochs)

    def get_current_ratio(self):
        """Get current ratio based on training progress"""
        if self.total_steps == 0:
            return self.initial_ratio
        
        progress = min(self.state.global_step / self.total_steps, 1.0)
        current_ratio = self.initial_ratio - (self.initial_ratio - self.final_ratio) * progress
        return max(int(round(current_ratio)), self.final_ratio)

    def get_train_dataloader(self):
        # Update dataset with total steps info
        if hasattr(self.train_dataset, 'total_steps'):
            self.train_dataset.total_steps = self.total_steps
            
        sampler = AlternatingDistributedSampler(
            self.train_dataset,
            num_replicas=torch.distributed.get_world_size(),
            rank=torch.distributed.get_rank(),
            shuffle=False,
        )

        return DataLoader(
            self.train_dataset,
            batch_size=self.args.per_device_train_batch_size,
            sampler=sampler,
            collate_fn=self.data_collator,
            drop_last=self.args.dataloader_drop_last,
            num_workers=0,
            pin_memory=self.args.dataloader_pin_memory,
        )

    def training_step(self, model, inputs, num_items_in_batch=None):
        # Update dataset with current step
        if hasattr(self.train_dataset, 'set_current_step'):
            self.train_dataset.set_current_step(self.state.global_step)
        
        return super().training_step(model, inputs, num_items_in_batch)

    def log(self, logs, start_time=None):
        super().log(logs, start_time)
        if self.is_world_process_zero():
            current_ratio = self.get_current_ratio()
            global_step = self.state.global_step
            
            # Log current ratio
            if "loss" in logs:
                wandb.log({"current_ratio": current_ratio, "global_step": global_step})
            
            # Each cycle is (current_ratio + 1) steps
            cycle_length = current_ratio + 1
            step_in_cycle = global_step % cycle_length
            
            # Only log to wandb if 'loss' is in the logs dictionary
            if "loss" in logs:
                if step_in_cycle < current_ratio:
                    # Text loss
                    wandb.log({"text_loss": logs["loss"], "text_step": self.text_step})
                    self.text_step += 1
                else:
                    # Audio loss
                    wandb.log({"audio_loss": logs["loss"], "audio_step": self.audio_step})
                    self.audio_step += 1

    def save_model(self, output_dir=None, _internal_call=False):
        if output_dir is None:
            output_dir = self.args.output_dir
        self.save_and_push_model(output_dir)

    def save_and_push_model(self, output_dir):
        save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
        with FSDP.state_dict_type(self.model, StateDictType.FULL_STATE_DICT, save_policy):
            cpu_state_dict = self.model.state_dict()
        self.model.save_pretrained(output_dir, state_dict=cpu_state_dict)


def data_collator(features):
    input_ids = [f["input_ids"] for f in features]

    if any("attention_mask" not in f for f in features):
        attention_mask = [[1]*len(ids) for ids in input_ids]
    else:
        attention_mask = [f["attention_mask"] for f in features]

    if any("labels" not in f for f in features):
        labels = input_ids
    else:
        labels = [f["labels"] for f in features]

    input_ids = torch.nn.utils.rnn.pad_sequence([torch.tensor(
        i, dtype=torch.long) for i in input_ids], batch_first=True, padding_value=pad_token)
    attention_mask = torch.nn.utils.rnn.pad_sequence([torch.tensor(
        m, dtype=torch.long) for m in attention_mask], batch_first=True, padding_value=0)
    labels = torch.nn.utils.rnn.pad_sequence([torch.tensor(
        l, dtype=torch.long) for l in labels], batch_first=True, padding_value=-100)

    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}


wandb.init(project=project_name, name=run_name)

# Setup accelerate (this initializes distributed environment)
accelerator = Accelerator()
device = accelerator.device

tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)

# Initialize model with proper dtype for Flash Attention 2.0 
model = AutoLigerKernelForCausalLM.from_pretrained(
    model_name, 
    attn_implementation="kernels-community/flash-attn3",
    torch_dtype=torch.bfloat16,
)

# Initialize model on first GPU to make Flash Attention happy
if accelerator.is_local_main_process:
    print(f"Pre-initializing model on {device} before FSDP")
model = model.to(device)

number_add_tokens = 7 * 4096 + 10
new_tokens = [f"<custom_token_{i}>" for i in range(0, number_add_tokens + 1)]
tokenizer.add_tokens(new_tokens)
model.resize_token_embeddings(len(tokenizer))

ds1 = load_dataset(dsn1, split="train")
ds2 = load_dataset(dsn2, split="train")

batch_total = batch_size * number_processes

# Calculate total steps for the dataset
num_update_steps_per_epoch = len(ds1) // (batch_size * number_processes)
total_steps = int(num_update_steps_per_epoch * epochs)

train_dataset = GradualRatioDataset(
    ds1, ds2, batch_total, 
    initial_ratio=initial_ratio, 
    final_ratio=final_ratio, 
    total_steps=total_steps
)

training_args = TrainingArguments(
    overwrite_output_dir=True,
    num_train_epochs=epochs,
    per_device_train_batch_size=batch_size,
    logging_steps=1,
    bf16=True,
    output_dir=f"./{base_repo_id}",
    fsdp="auto_wrap",
    report_to="wandb",
    save_steps=save_steps,
    remove_unused_columns=True,
    learning_rate=learning_rate,
    lr_scheduler_type="cosine", 
)

trainer = FSDPTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    data_collator=data_collator,
    initial_ratio=initial_ratio,
    final_ratio=final_ratio
)

print(f"Starting training with ratio progression: {initial_ratio}:1 -> {final_ratio}:1")
print(f"Total steps: {total_steps}")

trainer.train()
```
accelerate config: 
```bash
{
  "compute_environment": "LOCAL_MACHINE",
  "debug": false,
  "distributed_type": "MULTI_GPU",
  "downcast_bf16": false,
  "enable_cpu_affinity": true,
  "machine_rank": 0,
  "main_process_ip": "137.135.43.39",
  "main_process_port": 29500,
  "main_training_function": "main",
  "mixed_precision": "bf16",
  "num_machines": 2,
  "num_processes": 16,
  "rdzv_backend": "static",
  "same_network": true,
  "tpu_use_cluster": false,
  "tpu_use_sudo": false,
  "use_cpu": false,
  "fsdp_config": {
    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
    "fsdp_backward_prefetch_policy": "BACKWARD_PRE",
    "fsdp_forward_prefetch": true,
    "fsdp_offload_params": false,
    "fsdp_sharding_strategy": 1,
    "fsdp_state_dict_type": "FULL_STATE_DICT",
    "fsdp_transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"] # SmolLM3DecoderLayer, LlamaDecoderLayer
  },
}
```

### Expected behavior

I'm training a TTS based on LLaMA. It works with 8xH100. However, when I try to train in parallel with 8xH100 on two different servers, I encounter this error.

How can I resolve it?
```bash
rank0]: Traceback (most recent call last):
[rank0]:   File "/mnt/kadirnar/.../../train.py", line 285, in <module>
[rank0]:     training_args = TrainingArguments(
[rank0]:   File "<string>", line 133, in __init__
[rank0]:   File "/mnt/kadirnar/.../.venv/lib/python3.10/site-packages/transformers/training_args.py", line 1731, in __post_init__
[rank0]:     raise ValueError(error_message)
[rank0]: ValueError: Your setup doesn't support bf16/gpu.

```

Comments:
@ArthurZucker  I tried with 8xH100 and I'm getting the same error. I've trained a model a few times before. Could this error be related to something specific?

Hey! sorry we'll have a look! cc @matej-svejda as well! 

@kadirnar Could you send your environment from both of the nodes? Also can you try running this code on both of the nodes and tell me what you get, resp. check in which branch it fails?

```python
def is_torch_bf16_gpu_available() -> bool:
    if not is_torch_available():
        return False

    import torch

    if torch.cuda.is_available():
        return torch.cuda.is_bf16_supported()
    if is_torch_xpu_available():
        return torch.xpu.is_bf16_supported()
    if is_torch_hpu_available():
        return True
    if is_torch_npu_available():
        return torch.npu.is_bf16_supported()
    return False
```

> [@kadirnar](https://github.com/kadirnar) Could you send your environment from both of the nodes? Also can you try running this code on both of the nodes and tell me what you get, resp. check in which branch it fails?
> 
> def is_torch_bf16_gpu_available() -> bool:
>     if not is_torch_available():
>         return False
> 
>     import torch
> 
>     if torch.cuda.is_available():
>         return torch.cuda.is_bf16_supported()
>     if is_torch_xpu_available():
>         return torch.xpu.is_bf16_supported()
>     if is_torch_hpu_available():
>         return True
>     if is_torch_npu_available():
>         return torch.npu.is_bf16_supported()
>     return False

I just want to fix the 8xH100 error right now. I tried it on two different cloud services. The code worked on Hyperstack. However, it doesn't work on the Hyperbolid platform. Is it related to the Nvidia packages?

```bash
/mnt/data/.../.venv/lib/python3.10/site-packages/torch/cuda/__init__.py:182: UserWarning: CUDA initialization: Unexpected error from cudaGetDeviceCount(). Did you run some cuda functions before calling NumCudaDevices() that might have already set an error? Error 802: system not yet initialized (Triggered internally at /pytorch/c10/cuda/CUDAFunctions.cpp:109.)
  return torch._C._cuda_getDeviceCount() > 0
```

Yes this seems like system issue, not transformers, your cuda seems to not be correctly setup.

> Yes this seems like system issue, not transformers, your cuda seems to not be correctly setup.

I'm sorry for the simple mistake. Yes, it's an error related to cuda-toolkit.

@ArthurZucker 
I've tested model training on many different servers. And I have correctly installed cuda. Flash-attn2 and sdpa optimization model training are working. But it does work only in version 4.52. It gives an error in version 4.54.

```bash
[rank3]: Traceback (most recent call last):
[rank3]:   File "/mnt/kadirnar/github/.../.../train.py", line 312, in <module>
[rank3]:     trainer.train()
[rank3]:   File "/mnt/kadirnar/github/.../.venv/lib/python3.10/site-packages/transformers/trainer.py", line 2237, in train
[rank3]:     return inner_training_loop(
[rank3]:   File "/mnt/kadirnar/github/.../.venv/lib/python3.10/site-packages/transformers/trainer.py", line 2532, in _inner_training_loop
[rank3]:     batch_samples, num_items_in_batch = self.get_batch_samples(epoch_iterator, num_batches, args.device)
[rank3]:   File "/mnt/kadirnar/github/.../.venv/lib/python3.10/site-packages/transformers/trainer.py", line 5371, in get_batch_samples
[rank3]:     num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum()
[rank3]:   File "/mnt/kadirnar/github/.../.venv/lib/python3.10/site-packages/accelerate/accelerator.py", line 2794, in gather
[rank3]:     return gather(tensor)
[rank3]:   File "/mnt/kadirnar/github/.../.venv/lib/python3.10/site-packages/accelerate/utils/operations.py", line 371, in wrapper
[rank3]:     return function(*args, **kwargs)
[rank3]:   File "/mnt/kadirnar/github/.../.venv/lib/python3.10/site-packages/accelerate/utils/operations.py", line 432, in gather
[rank3]:     return _gpu_gather(tensor)
[rank3]:   File "/mnt/kadirnar/github/.../.venv/lib/python3.10/site-packages/accelerate/utils/operations.py", line 351, in _gpu_gather
[rank3]:     return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)
[rank3]:   File "/mnt/kadirnar/github/.../.venv/lib/python3.10/site-packages/accelerate/utils/operations.py", line 126, in recursively_apply
[rank3]:     return func(data, *args, **kwargs)
[rank3]:   File "/mnt/kadirnar/github/.../.venv/lib/python3.10/site-packages/accelerate/utils/operations.py", line 341, in _gpu_gather_one
[rank3]:     gather_op(output_tensors, tensor)
[rank3]:   File "/mnt/kadirnar/github/..../.venv/lib/python3.10/site-packages/torch/distributed/c10d_logger.py", line 81, in wrapper
[rank3]:     return func(*args, **kwargs)
[rank3]:   File "/mnt/kadirnar/github/.../.venv/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 3836, in all_gather_into_tensor
[rank3]:     work = group._allgather_base(output_tensor, input_tensor, opts)
[rank3]: ValueError: Tensors must be CUDA and dense
```

nvcc: Build cuda_12.1.r12.1/compiler.32688072_0


We recently changed the default value for `average_tokens_across_devices` to `True` and this might be the issue

```python
if self.args.average_tokens_across_devices:
                num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum() 
```

Can you check the value of `num_items_in_batch` across your processes so that we can find a solution to fix this ? 
Optionally, you can set it to `False` for now 

> We recently changed the default value for `average_tokens_across_devices` to `True` and this might be the issue
> 
> if self.args.average_tokens_across_devices:
>                 num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum() 
> Can you check the value of `num_items_in_batch` across your processes so that we can find a solution to fix this ? Optionally, you can set it to `False` for now

Yes, the error has been resolved, thank you very much. It would be very helpful to write an error message for this bug. 

@S1ro1 The problem I'm experiencing is related to accelerate. But can you answer in this issue? 

This is my first time training using more than 8xH100 GPUs. So, I might be making some simple mistakes. Training the model with 8xH100 is 40 hours faster than training with 16xH100.  When I increase the batch size value during model training with 16xH100, it gives a CUDA memory error. But I have more memory. Where am I making a mistake?

8xH100: 160 hour.         | 45/422676 [01:12<196:26:36,  1.67s/it
16xH100: 200 hour        | 40/211338 [02:21<195:20:57,  3.33s/it

8xh100(single node)
yaml:
```bash
{
  "compute_environment": "LOCAL_MACHINE",
  "debug": false,
  "distributed_type": "MULTI_GPU",
  "downcast_bf16": false,
  "enable_cpu_affinity": true,
  "machine_rank": 0,
  "main_training_function": "main",
  "mixed_precision": "bf16",
  "num_machines": 1,
  "num_processes": 8,
  "rdzv_backend": "static",
  "same_network": true,
  "tpu_use_cluster": false,
  "tpu_use_sudo": false,
  "use_cpu": false,
  "fsdp_config": {
    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
    "fsdp_backward_prefetch_policy": "BACKWARD_PRE",
    "fsdp_forward_prefetch": true,
    "fsdp_offload_params": false,
    "fsdp_sharding_strategy": 1,
    "fsdp_state_dict_type": "FULL_STATE_DICT",
    "fsdp_transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"] # SmolLM3DecoderLayer, LlamaDecoderLayer
  },
}
```

16xH100(2 node):
yaml:
```bash
{
  "compute_environment": "LOCAL_MACHINE",
  "debug": false,
  "distributed_type": "MULTI_GPU",
  "downcast_bf16": false,
  "enable_cpu_affinity": true,
  "machine_rank": 1,
  "main_training_function": "main",
  "mixed_precision": "bf16",
  "main_process_ip": "10.15.46.25",
  "main_process_port": 29500,
  "num_machines": 2,
  "num_processes": 16,
  "rdzv_backend": "static",
  "same_network": true,
  "tpu_use_cluster": false,
  "tpu_use_sudo": false,
  "use_cpu": false,
  "fsdp_config": {
    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
    "fsdp_backward_prefetch_policy": "BACKWARD_PRE",
    "fsdp_forward_prefetch": true,
    "fsdp_offload_params": false,
    "fsdp_sharding_strategy": 1,
    "fsdp_state_dict_type": "FULL_STATE_DICT",
    "fsdp_transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"] # SmolLM3DecoderLayer, LlamaDecoderLayer
  },
}
```


> 8xH100: 160 hour. | 45/422676 [01:12<196:26:36, 1.67s/it
> 16xH100: 200 hour | 40/211338 [02:21<195:20:57, 3.33s/it

The example you show tells that both are ~195h. Anyway I can't really help you without profiling output. I recommend profiling with torch compiler and try to figure out what's wrong there. I suppose that the communication is going over a slow interconnect, resulting in all-gathers in your forward to be a huge bottleneck.

Just got this error when upgraded transformers from `4.51.3` to `4.52.4`!
env: TPU v5e-8

lol, passed this error by upgrading transformers from `4.52.4` to `4.53.0`