| 220 | |
| 221 | |
| 222 | class Auto8bitTensor: |
| 223 | def __init__(self, data: Tensor, *args, **kwargs): |
| 224 | if isinstance(data, dict): # Add constructor from state dict |
| 225 | self._load_from_state_dict(data) |
| 226 | else: |
| 227 | abs_max = data.abs().max().item() |
| 228 | scale = abs_max / 127.0 if abs_max > 0 else 1.0 |
| 229 | |
| 230 | self.quantized = (data / scale).round().clamp(-127, 127).to(torch.int8) |
| 231 | self.scale = scale |
| 232 | self.orig_dtype = data.dtype |
| 233 | |
| 234 | def dequantize(self) -> Tensor: |
| 235 | return self.quantized.to(dtype=torch.float32) * self.scale |
| 236 | |
| 237 | def to(self, *args, **kwargs): |
| 238 | # Handle the dtype argument whether it's positional or keyword |
| 239 | dtype = None |
| 240 | if args and isinstance(args[0], torch.dtype): |
| 241 | dtype = args[0] |
| 242 | args = args[1:] |
| 243 | elif 'dtype' in kwargs: |
| 244 | dtype = kwargs['dtype'] |
| 245 | del kwargs['dtype'] |
| 246 | |
| 247 | if dtype is not None: |
| 248 | # First dequantize then convert to requested dtype |
| 249 | return self.dequantize().to(dtype=dtype, *args, **kwargs) |
| 250 | |
| 251 | # If no dtype specified, just pass through to parent |
| 252 | return self.dequantize().to(*args, **kwargs) |
| 253 | |
| 254 | def state_dict(self): |
| 255 | """Returns a dictionary containing the current state of the tensor.""" |
| 256 | return { |
| 257 | 'quantized': self.quantized, |
| 258 | 'scale': self.scale, |
| 259 | 'orig_dtype': self.orig_dtype |
| 260 | } |
| 261 | |
| 262 | def _load_from_state_dict(self, state_dict): |
| 263 | """Loads the tensor state from a state dictionary.""" |
| 264 | self.quantized = state_dict['quantized'] |
| 265 | self.scale = state_dict['scale'] |
| 266 | self.orig_dtype = state_dict['orig_dtype'] |
| 267 | |
| 268 | def __str__(self): |
| 269 | return f"Auto8bitTensor({self.dequantize()})" |
| 270 | |
| 271 | |
| 272 | def stochastic_grad_accummulation(param): |
no outgoing calls
no test coverage detected