Repository navigation
FSDP2: Full Activation Checkpointing does not work with torch.compile #361
Description
Apparently torch compile changes input and output shapes of the compiled regions, which breaks the expected shapes after recomputation of Activation Checkpointing in the backward pass. Torch compile and AC each work without any issues individually. Also, selective_op_activation_checkpointed_model has no issues when jointly used with torch.compile.
Error logs for torch.compile + full AC
[rank7]: torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed values for the following tensors have different metadata than during the forward pass.
[rank7]: tensor at position 2:
[rank7]: saved metadata: {'shape': torch.Size([4096]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=7)}
[rank7]: recomputed metadata: {'shape': torch.Size([1, 1, 4096, 128]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=7)}
[rank7]: tensor at position 3:
[rank7]: saved metadata: {'shape': torch.Size([1, 4096, 1]), 'dtype': torch.float32, 'device': device(type='cuda', index=7)}
[rank7]: recomputed metadata: {'shape': torch.Size([1, 1, 4096, 128]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=7)}
[rank7]: tensor at position 4:
[rank7]: saved metadata: {'shape': torch.Size([1, 4096, 1]), 'dtype': torch.float32, 'device': device(type='cuda', index=7)}
[rank7]: recomputed metadata: {'shape': torch.Size([4096]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=7)}
[rank7]: tensor at position 5:
[rank7]: saved metadata: {'shape': torch.Size([4096, 4096]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=7)}
[rank7]: recomputed metadata: {'shape': torch.Size([1, 4096, 1]), 'dtype': torch.float32, 'device': device(type='cuda', index=7)}
[rank7]: tensor at position 6:
[rank7]: saved metadata: {'shape': torch.Size([1, 1, 4096, 128]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=7)}
[rank7]: recomputed metadata: {'shape': torch.Size([1, 4096, 1]), 'dtype': torch.float32, 'device': device(type='cuda', index=7)}
[rank7]: tensor at position 7:
[rank7]: saved metadata: {'shape': torch.Size([1, 1, 4096, 128]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=7)}
[rank7]: recomputed metadata: {'shape': torch.Size([4096, 4096]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=7)}
respective config settings:
initialized_model:
component_key: model
variant_key: model_initialized
config:
model:
instance_key: fsdp_model
pass_type: BY_REFERENCE
model_initializer:
component_key: model_initialization
variant_key: composed
config:
model_type: gpt2
weight_init_type: scaled
mean: 0.0
std: 0.02
num_layers: ${model_raw.config.n_layer}
fsdp_model:
component_key: model
variant_key: fsdp2_wrapped
config:
model:
instance_key: compiled_model
pass_type: BY_REFERENCE
device_mesh:
instance_key: device_mesh
pass_type: BY_REFERENCE
mixed_precision_settings:
param_dtype: BF_16
reduce_dtype: BF_16
block_names: [GPT2Block]
compiled_model:
component_key: model
variant_key: compiled
config:
model:
instance_key: ${settings.benchmark.ac_mode}
pass_type: BY_REFERENCE
block_names: [GPT2Block]
full_activation_checkpointed_model:
component_key: model
variant_key: selective_activation_checkpointed
config:
sac_variant: full_activation_checkpointing
model:
instance_key: model_raw
pass_type: BY_REFERENCE
layers_fqn: transformer.h
Activity