muellerzr
commented
8 months ago cc @pacman100
Closed theoden8 closed 1 month ago
muellerzr
commented
8 months ago cc @pacman100
theoden8
commented
7 months ago Wrapping the model into a pytorch module doesn't help:
But it fails at unet call rather than vae.encode.
theoden8
commented
7 months ago If I use FSDP without accelerate it fails at loss.backward (which is further than accelerate), but the problem seems to be related:
github-actions[bot]
commented
6 months ago This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.
Please note that issues that do not follow the contributing guidelines are likely to be ignored.
theoden8
commented
6 months ago Any update on this? Could we try and reopen? @pacman100
github-actions[bot]
commented
5 months ago This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.
Please note that issues that do not follow the contributing guidelines are likely to be ignored.
theoden8
commented
4 months ago :/
samsara-ku
commented
2 months ago Is there any solution for this problem @theoden8 ? I stuck with this problem :(
theoden8
commented
2 months ago I stopped investigating some time ago
samsara-ku
commented
2 months ago Thanks for replying :+1:
drimeF0
commented
2 months ago same error, my code was supposed to train vae on 8 v2 TPU using fsdp
github-actions[bot]
commented
1 month ago This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.
Please note that issues that do not follow the contributing guidelines are likely to be ignored.
System Info
System info
```Shell Using 2 GPUs nvidia 2080ti # lsb_release -a Distributor ID: Ubuntu Description: Ubuntu 20.04.4 LTS Release: 20.04 Codename: focal # conda list | grep -E '(torch|transformers|accelerate|diffusers|deepspeed|numpy)' accelerate 0.25.0 pyhd8ed1ab_0 conda-forge deepspeed 0.12.2 cpu_py310h11dbdba_1 conda-forge diffusers 0.25.0 pyhd8ed1ab_0 conda-forge ffmpeg 4.3 hf484d3e_0 pytorch numpy 1.26.2 py310hb13e2d6_0 conda-forge pytorch 2.1.1 py3.10_cuda11.8_cudnn8.7.0_0 pytorch pytorch-cuda 11.8 h7e8668a_5 pytorch pytorch-lightning 2.1.2 pyhd8ed1ab_0 conda-forge pytorch-mutex 1.0 cuda pytorch torchaudio 2.1.1 py310_cu118 pytorch torchinfo 1.8.0 pyhd8ed1ab_0 conda-forge torchmetrics 1.2.1 pyhd8ed1ab_0 conda-forge torchtriton 2.1.0 py310 pytorch torchvision 0.16.1 py310_cu118 pytorch transformers 4.35.2 pyhd8ed1ab_0 conda-forge ```Information
Tasks
no_trainerscript in theexamplesfolder of thetransformersrepo (such asrun_no_trainer_glue.py)Reproduction
I am trying to train stable diffusion u-net with accelerate on multiple small GPUs. Here's my training script (adapted from https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet.py)
pipeline_accelerate_fsdp.py
```python #!/usr/bin/env python3 import os import sys import random import typing import numpy as np import tqdm.auto as tqdm import torch import torchvision import datasets import transformers import diffusers import accelerate import xformers import packaging if __name__ == '__main__': accelerator = accelerate.Accelerator() accelerator.print('fsdp', accelerator.local_process_index) sd_model = 'runwayml/stable-diffusion-v1-5' output_dir = 'output_ct_simple' pipeline = diffusers.StableDiffusionPipeline.from_pretrained( sd_model, torch_dtype=torch.float32, device_map=None, safety_checker=None, ) feature_extrator = pipeline.feature_extractor vae = pipeline.vae unet = pipeline.unet tokenizer = pipeline.tokenizer#.to(device=device) text_encoder = pipeline.text_encoder noise_scheduler = pipeline.scheduler vae.requires_grad_(True) unet.requires_grad_(False) text_encoder.requires_grad_(True) #tokenizer.requires_grad_(False) #print('after loading SDv1.5') #print_gpu_memory_usage() assert packaging.version.parse(xformers.__version__) != packaging.version.parse("0.0.16") unet.enable_xformers_memory_efficient_attention() dataset = datasets.load_dataset('fusing/fill50k', None, cache_dir=None) column_names = dataset['train'].column_names image_column = 'image' caption_column = 'text' conditioning_image_column = 'conditioning_image' def tokenize_captions(examples, is_train=True): proportion_empty_prompts = .0 captions = [] for caption in examples[caption_column]: if random.random() < proportion_empty_prompts: captions.append("") elif isinstance(caption, str): captions.append(caption) elif isinstance(caption, (list, np.ndarray)): # take a random caption if there are multiple captions.append(random.choice(caption) if is_train else caption[0]) else: raise ValueError( f"Caption column `{caption_column}` should contain either strings or lists of strings." ) inputs = tokenizer( captions, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt" ) return inputs.input_ids train_batch_size = 4 resolution = 128 image_transforms = torchvision.transforms.Compose([ torchvision.transforms.Resize(resolution, interpolation=torchvision.transforms.InterpolationMode.BILINEAR), torchvision.transforms.CenterCrop(resolution), torchvision.transforms.ToTensor(), torchvision.transforms.Normalize([0.5], [0.5]), ]) conditioning_image_transforms = torchvision.transforms.Compose([ torchvision.transforms.Resize(resolution, interpolation=torchvision.transforms.InterpolationMode.BILINEAR), torchvision.transforms.CenterCrop(resolution), torchvision.transforms.ToTensor(), ]) def preprocess_train(examples): images = [image.convert("RGB") for image in examples[image_column]] images = [image_transforms(image) for image in images] conditioning_images = [image.convert("RGB") for image in examples[conditioning_image_column]] conditioning_images = [conditioning_image_transforms(image) for image in conditioning_images] examples["pixel_values"] = images examples["conditioning_pixel_values"] = conditioning_images examples["input_ids"] = tokenize_captions(examples) return examples def collate_fn(examples): pixel_values = torch.stack([example["pixel_values"] for example in examples]) pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float() conditioning_pixel_values = torch.stack([example["conditioning_pixel_values"] for example in examples]) conditioning_pixel_values = conditioning_pixel_values.to(memory_format=torch.contiguous_format).float() input_ids = torch.stack([example["input_ids"] for example in examples]) return { "pixel_values": pixel_values, "conditioning_pixel_values": conditioning_pixel_values, "input_ids": input_ids, } train_dataset = dataset["train"].with_transform(preprocess_train) num_workers = 0 # main process train_dataloader = torch.utils.data.DataLoader( train_dataset, shuffle=True, collate_fn=collate_fn, batch_size=train_batch_size, num_workers=num_workers, ) learning_rate = 1e-4 trainable_params = [p for p in list(unet.parameters()) + list(vae.parameters()) + list(text_encoder.parameters()) if p.requires_grad] optimizer = torch.optim.Adam( trainable_params, lr=learning_rate, betas=(.9, .999), weight_decay=1e-2, eps=1e-8, ) lr_warmup_steps = 500 lr_num_cycles = 1 lr_power = 1. num_processes = 1 epochs = 10 lr_scheduler = diffusers.optimization.get_scheduler( 'constant', optimizer=optimizer, num_warmup_steps=lr_warmup_steps, num_training_steps=epochs * len(train_dataloader), num_cycles=lr_num_cycles, power=lr_power, ) unet, vae, text_encoder, tokenizer = accelerator.prepare( unet, vae, text_encoder, tokenizer ) optimizer, lr_scheduler, train_dataloader = accelerator.prepare( optimizer, lr_scheduler, train_dataloader ) if accelerator.is_main_process: pbar = tqdm.trange(epochs * len(train_dataset), initial=0, desc='steps', leave=True) for epoch in range(epochs): for batch in train_dataloader: pixel_values = batch["pixel_values"] pixel_values = pixel_values#.to(device=vae.device) latents = vae.encode(pixel_values).latent_dist.sample() latents = latents * vae.config.scaling_factor # Sample noise that we'll add to the latents noise = torch.randn_like(latents) bsz = latents.shape[0] # Sample a random timestep for each image timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device) timesteps = timesteps.long() # Add noise to the latents according to the noise magnitude at each timestep # (this is the forward diffusion process) noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps) # Get the text embedding for conditioning input_ids = batch['input_ids']#.to(device=text_encoder.device) encoder_hidden_states = text_encoder(input_ids)[0]#.float() pred = unet( noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states, ).sample if noise_scheduler.config.prediction_type == "epsilon": target = noise elif noise_scheduler.config.prediction_type == "v_prediction": target = noise_scheduler.get_velocity(latents, noise, timesteps) else: raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") loss = torch.nn.functional.mse_loss(pred, target, reduction="mean") optimizer.zero_grad() accelerator.backward(loss) optimizer.step() lr_scheduler.step() if accelerator.is_main_process: pbar.set_postfix(loss=loss.detach().cpu().item()) pbar.update(len(pixel_values)) ```accelerator_fsdp.yaml
```yaml compute_environment: LOCAL_MACHINE debug: false distributed_type: FSDP downcast_bf16: 'no' fsdp_config: fsdp_auto_wrap_policy: SIZE_BASED_WRAP fsdp_backward_prefetch_policy: BACKWARD_PRE fsdp_cpu_ram_efficient_loading: false fsdp_forward_prefetch: true fsdp_min_num_params: 100000000 fsdp_offload_params: false fsdp_sharding_strategy: 1 fsdp_state_dict_type: FULL_STATE_DICT fsdp_sync_module_states: true fsdp_use_orig_params: true machine_rank: 0 main_training_function: main mixed_precision: 'no' num_machines: 1 num_processes: 2 rdzv_backend: static same_network: true tpu_env: [] tpu_use_cluster: false tpu_use_sudo: false use_cpu: false ```run_pp_accelerate_fsdp.sh
```bash #!/usr/bin/env bash set -e eval "$("${HOME}"/.miniconda3/bin/conda shell.bash hook)" conda activate hfkernel set -x CUDA_HOME="/usr/local/cuda-11.1" \ OMP_NUM_THREADS=1 \ CUDA_VISIBLE_DEVICES=0,1 \ python -m accelerate.commands.launch \ --config_file ./accelerate_fsdp.yaml \ pipeline_accelerate_fsdp.py ```I get this error:
Expected behavior
I expected accelerator to launch FSDP with model parallelism and not pipeline parallelism, and shard models across available devices.
I understand that it flattens tensors in the underlying representation and shards it to different devices, so I'm trying to understand:
1.Why
nvidia-smishows that model loads to all devices, not sure if that's even sharded.I know there's deepspeed, but https://github.com/huggingface/accelerate/issues/253#issuecomment-1253231210 means I'd have to try and wrap everything into one model engine, so I was wondering if accelerate could handle it for me.