upgrade torchao to 0.17.0 (#3569)

* upgrade to torchao 0.17.0

* upgrade mistral-common too

* chore: lint

* patch fix for torchao low bit optimizers

* fix up

* propagate dtype

* fix test for ao change

* address PR comments

tests/e2e/test_quantization.py

@@ -5,20 +5,14 @@ Tests for axolotl.utils.quantization
 import pytest
 import torch
 from torch import nn
-from torchao.quantization import LinearActivationQuantizedTensor
+from torchao.quantization import IntxUnpackedToInt8Tensor
 from torchao.quantization.qat.embedding import FakeQuantizedEmbedding
 from torchao.quantization.qat.linear import FakeQuantizedLinear
 from torchao.quantization.quant_api import (
     Float8DynamicActivationFloat8WeightConfig,
     Float8DynamicActivationInt4WeightConfig,
+    Int8DynamicActivationIntxWeightConfig,
 )
-
-try:
-    from torchao.quantization.quant_api import Int8DynamicActivationInt4WeightConfig
-except ImportError:
-    from torchao.quantization.quant_api import (
-        Int8DynamicActivationIntxWeightConfig as Int8DynamicActivationInt4WeightConfig,
-    )
 from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor
 from transformers import AutoModelForCausalLM
 from transformers.trainer_callback import TrainerState
@@ -71,7 +65,7 @@ ptq_config_test_cases = [
         TorchAOQuantDType.int4,
         TorchAOQuantDType.int8,
         None,
-        Int8DynamicActivationInt4WeightConfig,
+        Int8DynamicActivationIntxWeightConfig,
     ),
     (
         TorchAOQuantDType.float8_e4m3fn,
@@ -96,7 +90,7 @@ ptq_test_cases = [
         8,
         False,
         None,
-        LinearActivationQuantizedTensor,
+        IntxUnpackedToInt8Tensor,
     ),
     # (
     #     TorchAOQuantDType.int4,