realign/train/configs/fsdp/llama2_fsdp_config.json

{
    "fsdp_transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"],
    "limit_all_gathers": true
}