No gradient update when using fsdp with hugginface accelerate

username: amarazad
cooked: <p>I am training a decoder LLM (GPT2 family). I converted code to accelerator following the instructions followed by huggingface.</p>
<p>When the accelerator config is without fsdp, it works fine. But when I enable fsdp in accelerator config, the gradient update is not happening, the model is static. Providing the non fsdp and fsdp config below.<br>
I have looked the code and searched blogs, issues but unable to find the issue.<br>
What could be the problem ? Any sugestions ?</p>
<p>environment variables:<br>
Python== 3.9<br>
transformers==4.28.1 <br>
accelerate==0.18.0 <br>
torch==1.13.1+cu117 --extra-index-url <a href="https://download.pytorch.org/whl/cu117" rel="noopener nofollow ugc">https://download.pytorch.org/whl/cu117</a> \</p>
<p>I assume FSDP enable code should also with single GPU.</p>
<p>Non-FSDP<br>
compute_environment: LOCAL_MACHINE<br>
distributed_type: ‘NO’<br>
downcast_bf16: ‘no’<br>
gpu_ids: all<br>
machine_rank: 0<br>
main_training_function: main<br>
mixed_precision: ‘no’<br>
num_machines: 1<br>
num_processes: 1<br>
rdzv_backend: static<br>
same_network: true<br>
tpu_env: []<br>
tpu_use_cluster: false<br>
tpu_use_sudo: false<br>
use_cpu: false</p>
<h2><a name="fsdp-config-1" class="anchor" href="#fsdp-config-1"></a>FSDP config</h2>
<p>compute_environment: LOCAL_MACHINE<br>
distributed_type: FSDP<br>
downcast_bf16: ‘no’<br>
fsdp_config:<br>
fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP<br>
fsdp_backward_prefetch_policy: BACKWARD_PRE<br>
fsdp_offload_params: false<br>
fsdp_sharding_strategy: 2<br>
fsdp_state_dict_type: FULL_STATE_DICT<br>
fsdp_transformer_layer_cls_to_wrap: CodeGenBlock<br>
machine_rank: 0<br>
main_training_function: main<br>
mixed_precision: ‘bf16’<br>
num_machines: 1<br>
num_processes: 1<br>
rdzv_backend: static<br>
same_network: true<br>
tpu_env: []<br>
tpu_use_cluster: false<br>
tpu_use_sudo: false<br>
use_cpu: false<br>
main_process_ip: x.x.x.x<br>
main_process_port: 53254</p>
<p>The key steps used in accelerator migration are:<br>
from accelerate import Accelerator<br>
from accelerate.utils import InitProcessGroupKwargs</p>
<p>from torch.distributed.fsdp import (<br>
FullyShardedDataParallel as FSDP,<br>
StateDictType,<br>
FullStateDictConfig<br>
)<br>
from accelerate import FullyShardedDataParallelPlugin<br>
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig</p>
<hr>
<h1><a name="initialize-2" class="anchor" href="#initialize-2"></a>initialize</h1>
<pre><code>launch_timeout = timedelta(seconds=config["ACCELERATOR"]["torch_launcher_to"])
accelerator_log_kwargs = {
        'kwargs_handlers':[InitProcessGroupKwargs(timeout=launch_timeout)],
        'gradient_accumulation_steps': config["ACCELERATOR"]['gradient_accumulation_steps'],
        }
accelerator_log_kwargs["log_with"] = "tensorboard"
accelerator_log_kwargs["project_dir"] = dirs["logs_dir"]

## Initialize the accelerator
accelerator = (
        Accelerator(**accelerator_log_kwargs)
        )
</code></pre>
<hr>
<p>model = accelerator.prepare(model) ### added for acclerator<br>
train_dataloader, optimizer, lr_scheduler, device = accelerator.prepare( train_dataloader, optimizer, lr_scheduler, device) ### added for acclerator</p>
<hr>
<pre><code>    loss = outputs.loss
    train_variables_dict['total_loss'] += loss.detach().float()
    optimizer.zero_grad()
    # loss.backward()
    # modified for accelerator
    accelerator.backward(loss) 
    optimizer.step()
    real_lr=optimizer.param_groups[0]['lr']
    lr_scheduler.step()
</code></pre>

username: ptrblck
cooked: <p>Could you update PyTorch to the latest stable (or better to the latest nightly) release and check if you would still see the same behavior?</p>