Support vectorized local reduction for p2p-based ReduceScatter overlap

[erhoo82]

Description

Vectorized load/store for p2p-based ReduceScatter overlap.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Vectorize the loads/stores of local reduction after TP ReduceScatter for (1) FP8 comm > FP32 reduce > BF16 out, (2) BF16 comm > FP32 reduce > BF16 out

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[timmoon10]

• Correctness bug when num_inputs == 1
• Unrolling loop over num_inputs is not necessary since it's not known at compile-time

Suggested change

	#pragma unroll
	for (int input_id = 1; input_id < num_inputs; ++input_id) {
	loader.load(tid + num_aligned_elements_per_input * input_id, tot_input_size);
	#pragma unroll
	for (int i = 0; i < nvec; ++i) {
	accum_buf[i] += static_cast<float>(loader.separate()[i]) * (*scale);
	if (input_id == num_inputs - 1) {
	storer.separate()[i] = static_cast<half_dtype>(accum_buf[i]);
	}
	}
	}
	for (int input_id = 1; input_id < num_inputs; ++input_id) {
	loader.load(tid + num_aligned_elements_per_input * input_id, tot_input_size);
	#pragma unroll
	for (int i = 0; i < nvec; ++i) {
	accum_buf[i] += static_cast<float>(loader.separate()[i]) * (*scale);
	}
	}
	#pragma unroll
	for (int i = 0; i < nvec; ++i) {
	storer.separate()[i] = static_cast<half_dtype>(accum_buf[i]);
	}

Same issue in reduce_bf16_cuda.

[timmoon10]

Since we're using this as a template arg, better to make it explicit that it is known at compile-time:

Suggested change

	const int nvec = 32;
	constexpr int nvec = 32;

Same issue in reduce_bf16.

[timmoon10]

Do we handle the case where the block size doesn't neatly divide the input size? Maybe we can fix with something like:

Suggested change

	const size_t tid = threadIdx.x + blockDim.x * blockIdx.x;
	const size_t tid = threadIdx.x + blockDim.x * blockIdx.x;
	if (tid >= num_aligned_elements_per_input) {
	return;
	}

Alternatively, we can change how we configure the CUDA blocks:

  size_t num_threads = MAX_THREADS / 4;
  assert(num_aligned_elements_per_input % num_threads == 0);
  size_t num_blocks = num_aligned_elements_per_input / num_threads;
  dim3 block(num_threads);
  dim3 grid(num_blocks);

Same issue in reduce_bf16_cuda.

[erhoo82]

@timmoon10
Can you review this one again?

BTW, regarding the lint error, nvec should be compile-time constant?