| |
| |
|
|
| """Host-only checks for greedy token selection with logprob reporting.""" |
|
|
| from types import SimpleNamespace |
|
|
| import pytest |
| import torch |
|
|
| import ttnn |
| from models.common.sampling.tt_sampling import TTSampling |
|
|
|
|
| @pytest.mark.parametrize("enabled", [False, True]) |
| def test_argmax_logprobs_use_original_shards_and_preserve_tokens(monkeypatch, enabled): |
| sampler = object.__new__(TTSampling) |
| sampler._force_argmax_sampling = True |
| sampler._force_argmax_sub_core_grids = None |
| sampler.mesh_device = SimpleNamespace(get_num_devices=lambda: 8) |
| sampler.tt_ccl = SimpleNamespace( |
| get_and_cycle_barrier_semaphore_handle=lambda axis: None, |
| get_and_cycle_ag_semaphore_handles=lambda axis: None, |
| ) |
| sampler.argmax_chunks_per_sync = 10 |
| sampler.argmax_num_workers_per_link = 1 |
| sampler._get_sampling_cluster_axis = lambda: 1 |
| sampler._get_force_argmax_all_gather_config = lambda axis: (1, ttnn.Topology.Linear) |
| sampler._can_slice_valid_vocab_for_argmax = lambda: False |
| sampler._mask_invalid_vocab_logits = lambda value: value |
| sampler._untilize_for_argmax = lambda value: value |
| shards = SimpleNamespace(memory_config=lambda: ttnn.DRAM_MEMORY_CONFIG) |
| gathered, probabilities = object(), object() |
| tokens = SimpleNamespace(shape=(1, 1, 32)) |
| canonical_tokens = object() |
| calls = [] |
|
|
| def calculate(logits, indices): |
| assert logits is shards, "Logprob normalization must not count replicated full vocabularies" |
| assert indices is canonical_tokens |
| calls.append(True) |
| return probabilities |
|
|
| def argmax(value, **kwargs): |
| assert value is gathered |
| return tokens |
|
|
| def reshape(value, shape): |
| assert value is tokens and shape == (1, 1, 1, 32) |
| return canonical_tokens |
|
|
| sampler.log_probs_calculator = SimpleNamespace(enable_log_probs=enabled, calculate_log_probs=calculate) |
| monkeypatch.setattr(ttnn.experimental, "all_gather_async", lambda *a, **kw: gathered) |
| monkeypatch.setattr(ttnn, "argmax", argmax) |
| monkeypatch.setattr(ttnn, "reshape", reshape) |
| result_tokens, result_probs = sampler.forward(shards) |
| assert result_tokens is tokens |
| assert result_probs is (probabilities if enabled else None) |
| assert len(calls) == int(enabled) |
|
|
|
|
| @pytest.mark.parametrize("enabled,use_topk,expected", [(False, True, True), (True, False, True), (True, True, False)]) |
| def test_topk_reporting_requires_candidate_pipeline(enabled, use_topk, expected): |
| sampler = object.__new__(TTSampling) |
| sampler._allow_force_argmax_sampling = True |
| sampler.log_probs_calculator = SimpleNamespace(enable_log_probs=enabled, _use_topk_logprobs=use_topk) |
| assert sampler._is_force_argmax_sampling([1], [1.0], [1.0]) is expected |
|
|
|
|
| def test_enabling_topk_reporting_refreshes_candidate_parameters(monkeypatch): |
| sampler = object.__new__(TTSampling) |
| sampler._allow_force_argmax_sampling = True |
| sampler.max_top_k = 32 |
| sampler._sampling_dp = 1 |
| sampler._greedy_col = object() |
| sampler.k_tensor, sampler.p_tensor, sampler.temp_tensor = object(), object(), object() |
| calculator = SimpleNamespace(enable_log_probs=False, _use_topk_logprobs=True) |
| calculator.set_log_probs_mode = lambda enabled, **kw: setattr(calculator, "enable_log_probs", enabled) |
| sampler.log_probs_calculator = calculator |
| copied = [] |
| monkeypatch.setattr(ttnn, "from_torch", lambda value, **kw: value) |
| monkeypatch.setattr(ttnn, "copy_host_to_device_tensor", lambda source, dest: copied.append((source, dest))) |
| sampler.reset_params(k=[1], p=[1.0], temp=[1.0], enable_log_probs=True) |
| assert sampler.force_argmax_sampling is False |
| assert len(copied) == 4 |
| assert copied[0][1] is sampler.k_tensor |
| assert torch.equal(copied[0][0], torch.tensor([1])) |
|
|
