Activation recomputation doc does not match implementation

[zinccat]

Describe the bug

The description in the doc (https://github.com/NVIDIA-NeMo/NeMo/blob/main/docs/source/features/optimizations/activation_recomputation.rst) on recompute_method does not match
https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/transformer_config.py#L290

[decimal-agent]

This is an AI-generated response. Please verify.

There appears to be a discrepancy in the activation recomputation feature between the NeMo documentation and the Megatron-LM implementation.

The NeMo documentation describes four recomputation methods:

• "full" - Checkpoints each transformer layer input
• "block" - Partial recomputation with specified number of layers
• "uniform" - Checkpoints inputs to blocks of consecutive layers
• "selective" - Self-attention recomputation without FlashAttention

However, there's an inconsistency in the documentation itself, where it uses "recompute_method" in most places but switches to "recompute_granularity" when describing the selective method:

recipe.model.config.recompute_granularity = "selective"  # Enable selective recomputation

This suggests the parameter name might be inconsistent between documentation and implementation, or there might be confusion between different parameters that control similar functionality.

Without being able to view the exact implementation code, I can't pinpoint the specific nature of the discrepancy, but it appears to involve either parameter naming, allowed values, or how these values are handled in the implementation.