Skip to content

PositionalEncoding re-uploads its whole table on every forward #19

Description

@CanReader

src/nn/embedding.rs:92:

pub struct PositionalEncoding { table: Tensor, dim: usize }
...
self.table.to(input.device()).index_select(&positions)

table is a bare Tensor, not a Buffer, so named_buffers() is empty and
Module::to_device never moves it. It sits on the host forever and every forward
copies max_len * dim floats host to device just to slice out the seq_len rows
it actually wants. At max_len = 4096, dim = 512 that's 8 MB per call.

Put the table in a Buffer and register it in named_buffers(), or cache the
device copy.

While in there, forward never asserts input.last_dim() == self.dim.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions