Support for multi tokens generation for vllm

[KpKqwq]

with model.generate(
prompt_input_vllm,
max_new_tokens=2048,
temperature=0.0,
top_p=1.0
) as generator:
outputs = model.generator.output.save()

this return no generate function

[Butanium]

@JadenFiotto-Kaufman I think it'd be nice for the vLLM class to have a wrapper around the generator to expose the same API as LanguageModel if it's not too hard. That's also something we can do in nnterp

[Butanium]

in the meantime @KpKqwq I think this snippet posted on discord shows how you can get the samples:

if __name__ == "__main__":
    from nnsight.modeling.vllm import VLLM

    model = VLLM("meta-llama/Llama-3.1-8B", dispatch=True, tensor_parallel_size=2)

    with model.trace(
        "The Eiffel Tower is located in the city of",
        temperatur=0.8,
        max_tokens=30,
    ) as tracer:

        activations = list().save()
        logits = list().save()
        samples = list().save()

        with tracer.iter[:30]:
            activations.append(model.model.layers[16].mlp.down_proj.output[0].cpu())
            logits = logits.append(model.logits.output)
            samples.append(model.samples.output)

        output = tracer.result.save()

[KpKqwq]

Thank, I will try that.