Tensor Parallelism Support for AffineQuantizedTensor

[jerryzh168]

Recently we landed #939 to support tensor parallelism for int8 weight only quantization, another example: #785

now we can support tensor parallelism for other types of quantization as well.

• float8 weight only @jainapurva - Add Float8 support for AQT tensor parallel #1003
• float8 dynamic activation @jainapurva - Add AQT tensor parallel for float8_dynamic_quant #1078
• uintx weight only @melvinebenezer
• int4 weight only quant - @jerryzh168 Add tensor parallelism support for int4_weight_only quantization #1120
• int8 dynamic act + int8 weight - not tested, add your name here
• fpx - add your name here

Steps

1. Create test

Since we don't have many tests today, we can optimize for readability for now, so we can copy paste the test cases to a https://github.com/pytorch/ao/blob/main/test/dtypes/test_affine_quantized_tensor_parallel.py instead of inheriting from these test cases

For new tests you can follow

ao/test/dtypes/test_affine_quantized_tensor_parallel.py

Lines 133 to 153 in c87cc9b

	class TestFloat8dqTensorAffineQuantizedTensorParallel(TestFloat8dqAffineQuantizedTensorParallel):
	QUANT_METHOD_FN = staticmethod(float8_dynamic_activation_float8_weight)
	QUANT_METHOD_KWARGS = {"granularity": PerTensor()}
	COMMON_DTYPES = [torch.bfloat16, torch.float16, torch.float32]
	@common_utils.parametrize("dtype", COMMON_DTYPES)
	@with_comms
	@unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
	def test_tp(self, dtype):
	return self._test_tp(dtype)
	class TestFloat8dqRowAffineQuantizedTensorParallel(TestFloat8dqAffineQuantizedTensorParallel):
	QUANT_METHOD_FN = staticmethod(float8_dynamic_activation_float8_weight)
	QUANT_METHOD_KWARGS = {"granularity": PerRow()}
	COMMON_DTYPES = [torch.bfloat16]
	@common_utils.parametrize("dtype", COMMON_DTYPES)
	@with_comms
	@unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
	def test_tp(self, dtype):
	return self._test_tp(dtype)

to create your own test case

2. Run the test

python test/dtypes/test_affine_quantized_tensor_parallel.py

3. Add support for missing ops until test passes

We'd expect people to add some slicing ops etc. to the corresponding TensorImpl tensor subclass

[melvinebenezer]

@jerryzh168 should'nt this Issue be kept open until other points in the description are fixed ?
I have a WIP #1026

[p4arth]

Hi @jerryzh168, I had two questions

1. What does fpx mean here exactly?
2. For int8 dynamic act + int8 weight only writing tests is left?

[jerryzh168]

@p4arth

• fpx is referring to this API:

  ao/torchao/quantization/quant_api.py

  Line 97 in 7aaf0ff

  "fpx_weight_only",

  , docs: https://github.com/pytorch/ao/tree/main/torchao/dtypes/floatx

• For int8 dynamic act + int8 weight only writing tests is left? - I think so, but will need to make sure the test passes, tests will be similar to

  ao/test/dtypes/test_affine_quantized_tensor_parallel.py

  Line 133 in 7aaf0ff

  class TestFloat8dqTensorAffineQuantizedTensorParallel(TestFloat8dqAffineQuantizedTensorParallel):