fsdp mixed precision pure vs default

`mixed_precision: PURE` and `mixed_precision: DEFAULT` in FSDP:

`mixed_precision: DEFAULT` (what you saw in logs):

- Parameters are stored in bfloat16

- Gradients are computed and reduced in float32

- Buffers (like batch norm stats) are in bfloat16

- Results in log: "param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.bfloat16"

`mixed_precision: PURE`:

- Parameters are stored in bfloat16

- Gradients are computed and reduced in bfloat16 (this is the key difference)

- Buffers are in bfloat16

- Would show in logs: "param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16, buffer_dtype=torch.bfloat16"

Performance comparison:

1. Memory Usage:

- PURE uses less memory because gradients are in bfloat16

- DEFAULT uses more memory because gradients are in float32

2. Speed:

- PURE is typically faster because:

  - Less memory bandwidth used during gradient communication

  - Faster gradient reduction operations

  - Particularly beneficial for distributed training

- However, training might be less stable

3. Training Stability:

- DEFAULT is more numerically stable because gradient reduction happens in float32

- PURE might require more careful tuning of learning rate and other hyperparameters

From your logs showing throughput around 191 tokens/sec/device, you might get better performance with PURE mode, possibly 5-15% faster due to reduced communication overhead. However, if you experience training instability (very high loss values or NaNs), you should switch back to DEFAULT.

Recommendation:

1. Start with PURE for better performance

2. Monitor training metrics closely

3. If you see instability, fall back to DEFAULT

hipblas, cublas algorithm

The HipBLASLt tuning process and algorithm selection is based on these factors in your data:

```

dev_cap,m,n,k,trans_a,trans_b,type_a,type_b,type_d,bias_type,lda,ldb,ldd,epi,comp,scale,ws_min,ws_max,algo_id,aidx

```

Key parameters:

1. Matrix Dimensions:

- `m,n,k`: Matrix dimensions for GEMM operations

- Example: `904,8192,2048,8192` = matrix sizes

2. Data Types:

- `type_a,type_b`: Input types (float8e4m3, bfloat16)

- `type_d`: Output type (bfloat16)

- `comp`: Computation type (f32)

3. Memory Layout:

- `trans_a,trans_b`: Matrix transposition (T=transposed, N=not)

- `lda,ldb,ldd`: Leading dimensions

4. Algorithm Selection:

- `algo_id`: Specific algorithm identifier

- `aidx`: Algorithm variant index

- workspace limits: `ws_min,ws_max`

The tuning process (`TE_HIPBLASLT_TUNING_RUN_COUNT=30` and `TE_HIPBLASLT_TUNING_ALGO_COUNT=100`) tests different combinations and selects the best based on:

1. Performance (speed)

2. Numerical stability

3. Memory usage

4. Hardware compatibility (dev_cap=904)

This tuning happens in the Tensor Engine (TE) library during the GEMM operations.

Insight about "fsdp_config.activation_checkpointing" option

 The `fsdp_config.activation_checkpointing` does come with a computational overhead since it recomputes activations during the backward pass, but it's generally the most efficient option for large models like LLaMA 70B for several reasons:

1. Alternative memory saving options usually have bigger throughput impacts:

- `activation_cpu_offload=true`: Much slower due to CPU-GPU data transfer

- `fp32_cpu_offload=true`: Also involves slow CPU-GPU transfers

- Reducing batch size: Directly reduces throughput

- Reducing model size: Changes model behavior

2. Activation checkpointing's overhead is typically around 20-30% compute time, but:

- Only affects backward pass

- Modern GPUs often have compute headroom

- Memory savings (50-70%) usually outweigh the computation cost

3. The recompute strategy is smart:

- Only recomputes necessary activations

- Can utilize GPU compute while other operations are happening

- Modern implementations are optimized to minimize recomputation

So while it's not zero-cost, it's usually the minimum-impact option that allows large models to fit in memory while maintaining reasonable throughput.

What is "GPU_MAX_HW_QUEUES"?

 GPU_MAX_HW_QUEUES controls the maximum number of hardware queues that can be used for GPU operations in ROCm/HIP environment (AMD GPUs).

Specifically:

1. Each hardware queue can handle GPU operations (kernels) independently
2. Default value is typically 8
3. In your script, it's set to 4: export GPU_MAX_HW_QUEUES=4

This setting affects:

• Parallel kernel execution
• Memory transfers
• Overall GPU utilization
• Resource allocation

Lower values (like 4) can:

• Reduce overhead
• Improve memory usage
• Lead to more predictable performance
• Be better for large models like LLaMA 70B

Higher values can:

• Allow more parallel operations
• Potentially improve throughput for smaller models
• Use more GPU resources

In your case, setting GPU_MAX_HW_QUEUES=4 is a conservative choice that helps manage memory and scheduling overhead when training the large LLaMA 70B model.

Hook Llama 3.1 8b layer and print dimension

 refer to code

.

def register_dimension_hooks(model, rank):

if rank != 0:

return

print('\n------------------- Model Structure -------------------')

print("Model type:", type(model))

# Get the actual model through the wrapper layers

if hasattr(model, 'model'):

model = model.model

if hasattr(model, 'model'):

model = model.model

print("Base model type:", type(model))

def make_hook(name, rank):

def hook(module, input, output):

print(f"\n--------------- Hook: {name} ---------------")

if hasattr(module, 'weight'):

weight = module.weight

print(f"GPU {rank} - {name}:")

print(f"Input shape: {input[0].shape}")

if hasattr(weight, '_local_tensor'):

local_weight = weight._local_tensor

print(f"Local weight shape: {local_weight.shape}")

print(f"Global weight shape: {weight.shape}")

if hasattr(weight, 'device_mesh'):

print(f"Device mesh: {weight.device_mesh}")

print(f"Placement: {weight.placements}")

print(f"Output shape: {output.shape}")

print("-" * 50)

return hook

# Register hooks for embedding layer

if hasattr(model, 'embed_tokens'):

print("Found embed_tokens")

model.embed_tokens.register_forward_hook(make_hook('embed_tokens', rank))

# Register hooks for all transformer layers

if hasattr(model, 'layers'):

for i, layer in enumerate(model.layers):

# Attention blocks

layer.self_attn.q_proj.register_forward_hook(

make_hook(f'layer_{i}_q_proj', rank))

layer.self_attn.k_proj.register_forward_hook(

make_hook(f'layer_{i}_k_proj', rank))

layer.self_attn.v_proj.register_forward_hook(

make_hook(f'layer_{i}_v_proj', rank))

layer.self_attn.o_proj.register_forward_hook(

make_hook(f'layer_{i}_o_proj', rank))

# MLP blocks

layer.mlp.gate_proj.register_forward_hook(

make_hook(f'layer_{i}_mlp_gate_proj', rank))

layer.mlp.up_proj.register_forward_hook(

make_hook(f'layer_{i}_mlp_up_proj', rank))

layer.mlp.down_proj.register_forward_hook(

make_hook(f'layer_{i}_mlp_down_proj', rank))

# Layer norms

layer.input_layernorm.register_forward_hook(

make_hook(f'layer_{i}_input_layernorm', rank))

layer.post_attention_layernorm.register_forward_hook(

make_hook(f'layer_{i}_post_attention_layernorm', rank))

# Register hook for final layer norm

if hasattr(model, 'norm'):

model.norm.register_forward_hook(make_hook('final_layernorm', rank))

# Register hook for LM head

if hasattr(model, 'lm_head'):

print("Found lm_head")

model.lm_head.register_forward_hook(make_hook('lm_head', rank))

# Print model structure to debug

print("\nModel attributes:", dir(model))

..

Thank you.

How FSDP, TP works on distributed training?

 5 different case of FSDP and TP usage.

FSDP and TP explanation for 2 layer model

 FSDP and TP are complementary parallelism techniques:

1. FSDP (Fully Sharded Data Parallelism):
   • Shards model parameters across GPUs
   • Each GPU holds a portion of each layer's parameters
   • During forward/backward pass, it gathers/scatters parameters as needed
   • Reduces memory usage per GPU, allowing larger models
2. TP (Tensor Parallelism):
   • Splits individual tensors (layers) across GPUs
   • Each GPU computes a portion of a layer's operations
   • Useful for very large layers that don't fit on a single GPU

When combined:

• FSDP handles overall model distribution
• TP handles distribution of large individual layers
• This allows for even larger models and better GPU utilization

Textual Representation:

   GPU 1        GPU 2        GPU 3        GPU 4
 +--------+   +--------+   +--------+   +--------+
 | L1 P1  |   | L1 P2  |   | L2 P1  |   | L2 P2  |
 |  TP1   |   |  TP2   |   |  TP1   |   |  TP2   |
 +--------+   +--------+   +--------+   +--------+
     |            |            |            |
     +------------+            +------------+
          Layer 1                  Layer 2

L1, L2: Layers 1 and 2
P1, P2: Parameter shards (FSDP)
TP1, TP2: Tensor Parallel splits

AMD Distributed Training Overview

# AMD Distributed Training Overview

AMD's approach to distributed training leverages its high-performance CPUs and GPUs, along with software frameworks, to enable efficient scaling of machine learning workloads across multiple nodes. Key aspects include:

1. **Hardware Solutions:**

   - AMD EPYC CPUs: Provide high core counts and memory bandwidth.

   - AMD Instinct GPUs: Accelerators designed for HPC and AI workloads.

   - AMD Infinity Fabric: High-speed interconnect for multi-GPU and multi-node systems.

2. **Software Framework:**

   - ROCm (Radeon Open Compute): Open-source software stack for GPU computing.

   - HIP (Heterogeneous-Compute Interface for Portability): C++ runtime API for GPU programming.

   - AMD's optimized libraries for deep learning frameworks like TensorFlow and PyTorch.

3. **Distributed Training Techniques:**

   - Data Parallelism: Distributing batches of training data across multiple GPUs or nodes.

   - Model Parallelism: Splitting large models across multiple devices.

   - Pipeline Parallelism: Dividing model layers across devices and processing in a pipelined fashion.

4. **Communication Optimization:**

   - RCCL (ROCm Communication Collectives Library): Optimized multi-GPU and multi-node collective communications.

   - Support for high-speed networking technologies like InfiniBand.

5. **Scalability:**

   - Support for scaling from single-node multi-GPU systems to large clusters.

   - Integration with job schedulers and resource managers for cluster environments.

6. **Ecosystem Integration:**

   - Compatibility with popular ML frameworks and distributed training tools.

   - Support for containers and orchestration platforms like Docker and Kubernetes.

7. **Performance Optimization:**

   - Mixed-precision training support.

   - Memory management techniques for large model training.

   - Automatic performance tuning tools.

AMD's distributed training solutions aim to provide high performance, scalability, and ease of use for researchers and organizations working on large-scale machine learning projects.