How to Perform Generation with Temperature, Top-K, and Top-P Sampling in Nanotron

You can perform generation with temperature, top_k, and top_p sampling in Nanotron by configuring GenerationArgs in src/nanotron/config/config.py and passing it to the decode_text function or run_generate.py CLI, which instantiates TopKSampler or TopPSampler from src/nanotron/generation/sampler.py to apply temperature scaling and probability filtering.

Nanotron provides a distributed generation pipeline that supports configurable sampling strategies for large language model inference. The framework handles tensor-parallel gathering of logits internally while exposing high-level APIs for temperature, top_k, and top_p constraints. Understanding the three-tier architecture—configuration, decoder engine, and sampler implementations—allows you to control randomness and diversity in generated text across multi-GPU setups.

Configuring Generation Parameters

The entry point for all sampling behavior is the GenerationArgs class defined in src/nanotron/config/config.py. This dataclass exposes the parameters that control how tokens are selected during generation:

• sampler: Specifies the strategy as a SamplerType enum value (greedy, top_k, top_p, basic)
• temperature: A float value (typically 0.0 to 2.0) that scales logits before softmax computation
• top_k: An integer limiting sampling to the k highest-probability tokens
• top_p: A float between 0.0 and 1.0 enabling nucleus sampling, selecting the smallest set of tokens whose cumulative probability exceeds the threshold

When you instantiate GenerationArgs with non-default values, the decode engine in src/nanotron/generation/decode.py routes the configuration to the appropriate sampler class during the forward pass.

Sampler Implementations

The core logic for temperature scaling and probability filtering resides in src/nanotron/generation/sampler.py. Each sampler operates on sharded logits gathered across tensor-parallel ranks via all-to-all communication, ensuring the full vocabulary is visible despite distributed model parallelism.

Temperature Scaling

Temperature scaling occurs immediately before the softmax operation to control the sharpness of the probability distribution. In TopPSampler, the division happens on line 69:

logits = logits / self.temperature

Similarly, TopKSampler applies temperature on line 225 before computing probabilities over the filtered token set. Lower temperatures (< 1.0) make the model more deterministic by sharpening the distribution, while higher temperatures (> 1.0) increase randomness.

Top-K Sampling

The TopKSampler class filters logits to retain only the k highest values per shard, gathers these across tensor-parallel ranks, applies temperature scaling, and computes a softmax over the reduced vocabulary subset. The actual sampling occurs on line 28 after the probability distribution is normalized. This method prevents the model from selecting low-probability outliers while maintaining diversity among high-likelihood candidates.

Top-P (Nucleus) Sampling

The TopPSampler implements nucleus sampling by sorting tokens by probability and accumulating their mass until reaching the threshold p (the self.p field). Lines 70-84 handle the masking of tokens outside the nucleus and sample from the truncated distribution. This dynamic approach adapts the vocabulary size per step based on the model's confidence, often producing more coherent long-form text than fixed top_k constraints.

Implementation Methods

You can trigger temperature-scaled sampling through three interfaces depending on your use case.

Using the decode_text Helper

For most applications, use the high-level decode_text function from src/nanotron/generation/decode.py:

from nanotron.generation.decode import decode_text
from nanotron.generation.sampler import SamplerType
from nanotron.config.config import GenerationArgs, TokenizerConfig
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
prompt = "Explain the concept of attention in transformers in two sentences."

gen_cfg = GenerationArgs(
    sampler=SamplerType.TOP_K,
    temperature=0.7,
    top_k=50,
)

outputs = decode_text(
    input_iter=(prompt,),
    tokenizer=tokenizer,
    model=model,
    parallel_context=parallel_context,
    max_new_tokens=100,
    generation_config=gen_cfg,
    tokenizer_config=TokenizerConfig(max_input_length=None),
)

for out in outputs:
    print(tokenizer.decode(out.generation_ids[0], skip_special_tokens=True))

The sampler parameter selects the TopKSampler class, which applies the 0.7 temperature scale during generation as implemented on line 225 of sampler.py.

Using the CLI

For quick testing without writing Python scripts, use run_generate.py with command-line flags:

python run_generate.py \
  --checkpoint_path=/path/to/checkpoint \
  --tokenizer_path=mistralai/Mistral-7B-v0.1 \
  --max_new_tokens=150 \
  --use_cache \
  --generation_sampler=top_p \
  --generation_temperature=0.9 \
  --generation_top_p=0.92

This script constructs a GenerationArgs instance from the flags (see lines 80-90 in run_generate.py) and forwards it to the same distributed decode engine used by the Python API.

Direct Sampler Invocation

For custom generation loops or research experiments, instantiate samplers directly:

from nanotron.generation.sampler import TopKSampler, TopPSampler

# Top-K with temperature

sampler = TopKSampler(pg=tp_pg, k=40, temperature=0.6)
next_token_ids = sampler(sharded_logits)  # Returns [batch] tensor

# Top-P with temperature

sampler = TopPSampler(pg=tp_pg, p=0.95, temperature=0.8)
next_token_ids = sampler(sharded_logits)

These calls bypass the higher-level decode logic but maintain the same distributed semantics, performing all-to-all gathers and temperature scaling exactly as implemented in src/nanotron/generation/sampler.py.

Summary

• GenerationArgs in src/nanotron/config/config.py specifies temperature, top_k, and top_p parameters for the generation pipeline.
• The decode engine in src/nanotron/generation/decode.py routes configurations to TopKSampler, TopPSampler, GreedySampler, or BasicSampler based on the sampler field.
• Temperature scaling divides logits by the temperature value before softmax in both TopKSampler (line 225) and TopPSampler (line 69).
• Top-K sampling restricts selection to the k highest-probability tokens gathered across tensor-parallel ranks.
• Top-P sampling dynamically selects the nucleus of tokens comprising cumulative probability p, implemented in lines 70-84 of sampler.py.
• You can access these features via the decode_text Python API, the run_generate.py CLI, or direct sampler instantiation.

Frequently Asked Questions

How does temperature affect token selection in Nanotron?

Temperature scales the logits by dividing by the temperature value before applying softmax. According to the source code in src/nanotron/generation/sampler.py, values below 1.0 increase confidence in high-probability tokens by sharpening the distribution, while values above 1.0 flatten the distribution to encourage diversity. The scaling occurs in both TopPSampler (line 69) and TopKSampler (line 225).

What is the difference between top_k and top_p sampling?

Top_k sampling restricts the candidate pool to a fixed number of highest-probability tokens regardless of their absolute probabilities, while top_p (nucleus) sampling selects a dynamic subset of tokens whose cumulative probability exceeds the threshold p. In Nanotron's implementation, TopKSampler filters by rank order before gathering across tensor-parallel groups, whereas TopPSampler calculates cumulative probabilities on the full vocabulary to determine the cutoff point.

Where is the sampling logic implemented in the codebase?

The sampling logic is implemented in src/nanotron/generation/sampler.py, which contains TopKSampler, TopPSampler, GreedySampler, and BasicSampler. The selection and instantiation of these samplers occurs in src/nanotron/generation/decode.py based on the GenerationArgs configuration defined in src/nanotron/config/config.py.

Can I use multiple sampling strategies simultaneously?

No, Nanotron's current architecture requires selecting a single sampler type via the sampler field in GenerationArgs. You cannot combine top_k truncation with top_p nucleus filtering in a single generation call; you must choose either SamplerType.TOP_K, SamplerType.TOP_P, or another supported type. To achieve combined effects, you would need to implement a custom sampler class following the pattern established in src/nanotron/generation/sampler.py.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}