(tllm_inter: Dict, hf_ref: Tuple[torch.Tensor], attn_mask,
logger)
| 178 | |
| 179 | |
| 180 | def intermediate_check(tllm_inter: Dict, hf_ref: Tuple[torch.Tensor], attn_mask, |
| 181 | logger): |
| 182 | |
| 183 | def apply_mask(x): |
| 184 | return x * attn_mask |
| 185 | |
| 186 | # minus one because there is an embedding output |
| 187 | num_layers = len(hf_ref) - 1 |
| 188 | |
| 189 | res = tllm_inter['embedding_output'] |
| 190 | res = apply_mask(res) |
| 191 | ref = hf_ref[0] |
| 192 | ref = apply_mask(ref) |
| 193 | torch.testing.assert_close(actual=res, expected=ref, rtol=1e-2, atol=1e-2) |
| 194 | logger.debug("Embedding are all close") |
| 195 | |
| 196 | for i in range(num_layers - 1): |
| 197 | res = tllm_inter[f'layer_{i}_output'] |
| 198 | res = apply_mask(res) |
| 199 | ref = hf_ref[i + 1] |
| 200 | ref = apply_mask(ref) |
| 201 | is_close = torch.allclose(res, ref, rtol=1e-2, atol=1e-2) |
| 202 | logger.debug(f'BertEncoderLayer_{i}_output is close: {is_close}') |
| 203 | |
| 204 | |
| 205 | @contextmanager |
no test coverage detected