nf4_detach missing args error when calling fully_shard on latest nightly

[ebsmothers]

Hi, I am seeing the error TypeError: nf4_detach() missing 1 required positional argument: 'args' when running torchtune's QLoRA FSDP2 recipe on the latest ao nightlies. Full stack trace

Repro:

tune download meta-llama/Llama-2-7b-hf --output-dir /tmp/Llama-2-7b-hf --hf-token <HF_TOKEN>
tune run --nnodes 1 --nproc_per_node 4 lora_finetune_fsdp2 --config llama2/7B_qlora enable_activation_checkpointing=False

Looks like it's happening when we try to call fully_shard on a TransformerDecoderLayer. When I downgrade to e.g. the 6/8 nightly things seem to work.

[msaroufim]

This might have been caused by #360 - cc @jerryzh168
@weifengpy did we ever end up merging some sanity tests in ao to detect these kinds of breakages cause relying purely on tune for testing seems less than ideal?

[weifengpy]

good catch! @ebsmothers

@msaroufim you are right. I will also check FSDP e2e in TorchAO just like how we do it in TorchTune

This might have been caused by #360 - cc @jerryzh168 @weifengpy did we ever end up merging some sanity tests in ao to detect these kinds of breakages cause relying purely on tune for testing seems less than ideal?

[supriyar]

@weifengpy can we close this issue now?

[weifengpy]

@weifengpy can we close this issue now?

Yes, we can close it as it's landed